foxhunt/crates/ml/tests/ensemble_inference_integration_test.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
    unused_crate_dependencies,
)]
//! Integration test: 4-model inference ensemble (DQN + PPO + Mamba2 + TFT)
//!
//! Constructs all four adapters with small configs, pre-warms the
//! sequence-based models (Mamba2, TFT) by feeding feature vectors into
//! their buffers, then runs the full ensemble and verifies bounded
//! predictions from all 4 models.

use ml::dqn::DQNConfig;
use ml::ensemble::adapters::{
    DqnInferenceAdapter, Mamba2InferenceAdapter, PpoInferenceAdapter, TftInferenceAdapter,
};
use ml::ensemble::inference_ensemble::InferenceEnsemble;
use ml::ensemble::{FeatureVector, ModelInferenceAdapter};
use ml::mamba::Mamba2Config;
use ml::ppo::PPOConfig;
use ml::tft::TFTConfig;

/// Small DQN config -- 51 input features, 45 actions, two small hidden layers.
fn small_dqn_config() -> DQNConfig {
    DQNConfig {
        num_actions: 63,
        hidden_dims: vec![64, 64],
        ..Default::default()
    }
}

/// Small PPO config -- 64 input (zero-padded from 51), 45 actions, small MLPs.
fn small_ppo_config() -> PPOConfig {
    PPOConfig {
        state_dim: 64,
        num_actions: 63,
        policy_hidden_dims: vec![64, 64],
        value_hidden_dims: vec![64, 64],
        ..Default::default()
    }
}

/// Small Mamba2 config -- d_model=64 (>51 features), 1 layer, seq_len=4.
fn small_mamba2_config() -> Mamba2Config {
    Mamba2Config {
        d_model: 64,
        d_state: 16,
        d_head: 16,
        num_heads: 2,
        expand: 2,
        num_layers: 1,
        max_seq_len: 8,
        dropout: 0.0,
        ..Default::default()
    }
}

/// Small TFT config -- input_dim=20, 1 layer, seq_len=4.
/// input_dim = num_static(6) + num_known(6) + num_unknown(8) = 20
fn small_tft_config() -> TFTConfig {
    TFTConfig {
        input_dim: 20,
        hidden_dim: 32,
        num_heads: 2,
        num_layers: 1,
        prediction_horizon: 5,
        sequence_length: 4,
        num_quantiles: 9,
        num_static_features: 6,
        num_known_features: 6,
        num_unknown_features: 8,
        dropout_rate: 0.0,
        ..Default::default()
    }
}

use tracing::info;

const SEQ_LEN: usize = 4;
const NUM_FEATURE_VECTORS: usize = 5;

/// Build a 51-dimensional feature vector with a distinguishing pattern.
fn make_feature_vector(index: usize) -> FeatureVector {
    let base = 0.01 * (index as f64 + 1.0);
    FeatureVector {
        values: vec![base; 51],
        timestamp: 1_700_000_000_000_000 + index as i64,
    }
}

#[test]
fn test_four_model_ensemble_integration() {
    // ---- construct all 4 adapters ----
    let dqn = DqnInferenceAdapter::new(small_dqn_config())
        .expect("DqnInferenceAdapter::new should succeed with small config");

    let ppo = PpoInferenceAdapter::new(small_ppo_config())
        .expect("PpoInferenceAdapter::new should succeed with small config");

    let mamba2 = Mamba2InferenceAdapter::new(small_mamba2_config(), SEQ_LEN)
        .expect("Mamba2InferenceAdapter::new should succeed with small config");

    let tft = TftInferenceAdapter::new(small_tft_config(), SEQ_LEN)
        .expect("TftInferenceAdapter::new should succeed with small config");

    // ---- DQN + PPO are ready immediately; Mamba2 + TFT need buffer fill ----
    assert!(dqn.is_ready(), "DQN should be ready immediately");
    assert!(ppo.is_ready(), "PPO should be ready immediately");
    assert!(
        !mamba2.is_ready(),
        "Mamba2 should NOT be ready (empty buffer)"
    );
    assert!(!tft.is_ready(), "TFT should NOT be ready (empty buffer)");

    // ---- pre-warm sequence-based adapters (Mamba2, TFT) ----
    // The ensemble only calls predict on ready adapters, so we must fill
    // the internal sequence buffers before assembling the ensemble.
    for i in 0..SEQ_LEN {
        let fv = make_feature_vector(i);
        let mamba2_pred = mamba2
            .predict(&fv)
            .unwrap_or_else(|e| panic!("Mamba2 warm-up {} should not error: {}", i, e));
        let tft_pred = tft
            .predict(&fv)
            .unwrap_or_else(|e| panic!("TFT warm-up {} should not error: {}", i, e));

        // Before the buffer is full, direction and confidence should be neutral (0.0)
        if i < SEQ_LEN - 1 {
            assert_eq!(
                mamba2_pred.direction, 0.0,
                "Mamba2 should return neutral direction during warm-up (step {})",
                i
            );
            assert_eq!(
                tft_pred.direction, 0.0,
                "TFT should return neutral direction during warm-up (step {})",
                i
            );
        } else {
            // Buffer is now full -- should produce real predictions
            assert!(
                mamba2_pred.direction >= -1.0 && mamba2_pred.direction <= 1.0,
                "Mamba2 warm-up final: direction {} out of [-1, 1]",
                mamba2_pred.direction
            );
            assert!(
                tft_pred.direction >= -1.0 && tft_pred.direction <= 1.0,
                "TFT warm-up final: direction {} out of [-1, 1]",
                tft_pred.direction
            );
        }

        info!(
            i,
            mamba2_direction = mamba2_pred.direction,
            mamba2_confidence = mamba2_pred.confidence,
            tft_direction = tft_pred.direction,
            tft_confidence = tft_pred.confidence,
            "Warm-up step"
        );
    }

    // All 4 adapters should now be ready
    assert!(mamba2.is_ready(), "Mamba2 should be ready after warm-up");
    assert!(tft.is_ready(), "TFT should be ready after warm-up");

    // ---- build the ensemble with all 4 pre-warmed adapters ----
    let adapters: Vec<Box<dyn ModelInferenceAdapter>> =
        vec![Box::new(dqn), Box::new(ppo), Box::new(mamba2), Box::new(tft)];
    let ensemble = InferenceEnsemble::new(adapters);

    // All 4 models should be ready
    assert_eq!(
        ensemble.ready_count(),
        4,
        "All 4 adapters should be ready after warm-up"
    );

    // ---- feed 5 feature vectors through the full ensemble ----
    for i in 0..NUM_FEATURE_VECTORS {
        let fv = make_feature_vector(SEQ_LEN + i);
        let pred = ensemble
            .predict(&fv)
            .unwrap_or_else(|e| panic!("Ensemble prediction {} should succeed: {}", i, e));

        // Direction must be in [-1.0, 1.0]
        assert!(
            pred.direction >= -1.0 && pred.direction <= 1.0,
            "FV {}: direction {} out of [-1, 1]",
            i,
            pred.direction
        );

        // Confidence must be in [0.0, 1.0]
        assert!(
            pred.confidence >= 0.0 && pred.confidence <= 1.0,
            "FV {}: confidence {} out of [0, 1]",
            i,
            pred.confidence
        );

        // All 4 models should contribute (ensemble name includes all)
        assert!(
            pred.model_name.contains("DQN"),
            "FV {}: ensemble name should include DQN, got: {}",
            i,
            pred.model_name
        );
        assert!(
            pred.model_name.contains("PPO"),
            "FV {}: ensemble name should include PPO, got: {}",
            i,
            pred.model_name
        );
        assert!(
            pred.model_name.contains("MAMBA-2"),
            "FV {}: ensemble name should include MAMBA-2, got: {}",
            i,
            pred.model_name
        );
        assert!(
            pred.model_name.contains("TFT"),
            "FV {}: ensemble name should include TFT, got: {}",
            i,
            pred.model_name
        );

        info!(
            i,
            direction = pred.direction,
            confidence = pred.confidence,
            model_name = %pred.model_name,
            "Ensemble prediction"
        );
    }

    // ---- confirm ready_count is still 4 after all predictions ----
    assert_eq!(
        ensemble.ready_count(),
        4,
        "All 4 adapters should remain ready after {} ensemble predictions",
        NUM_FEATURE_VECTORS
    );

    info!("All 4-model ensemble integration checks passed");
}