foxhunt/testing/integration/integration/checkpoint_roundtrip.rs

//! Checkpoint roundtrip integration tests
//!
//! Verifies that saving a model checkpoint and loading it into a fresh model
//! produces identical inference output. Tests cover:
//!
//! - DQN: raw Q-network weight save/load via safetensors
//! - DQN: adapter roundtrip via `DqnInferenceAdapter::from_checkpoint`
//! - DQN: `CheckpointManager` + `Checkpointable` trait flow
//! - PPO: checkpoint save/load roundtrip
//! - TFT: dual adapter deterministic inference with sequence buffer
//! - TFT: quantile metadata verification after buffer fill
//! - Mamba2: dual adapter deterministic inference with sequence buffer
//! - Mamba2: `CheckpointManager` + `Checkpointable` trait on `Mamba2SSM`

use std::sync::OnceLock;

use ml::checkpoint::{CheckpointConfig, CheckpointManager, CompressionType};
use ml::ensemble::adapters::{
    Mamba2InferenceAdapter, PpoInferenceAdapter, TftInferenceAdapter,
};
use ml::ensemble::inference_adapter::{FeatureVector, ModelInferenceAdapter};
use ml::mamba::{Mamba2Config, Mamba2SSM};
use ml::ppo::ppo::{PPOConfig, PPO};
use ml::prelude::MlDevice;
use ml::tft::TFTConfig;

// ---------------------------------------------------------------------------
// Shared CUDA device
// ---------------------------------------------------------------------------

static SHARED_CUDA: OnceLock<MlDevice> = OnceLock::new();

fn cuda_device() -> MlDevice {
    SHARED_CUDA
        .get_or_init(|| MlDevice::cuda(0).expect("CUDA device required"))
        .clone()
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn small_ppo_config() -> PPOConfig {
    PPOConfig {
        state_dim: 64,
        num_actions: 45,
        policy_hidden_dims: vec![32, 32],
        value_hidden_dims: vec![32, 32],
        ..Default::default()
    }
}

fn fixed_feature_vector_64() -> FeatureVector {
    FeatureVector {
        values: vec![0.3; 64],
        timestamp: 1_700_000_000,
    }
}

// ---------------------------------------------------------------------------
// PPO: checkpoint roundtrip
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_ppo_checkpoint_roundtrip() {
    let config = small_ppo_config();

    // 1. Create PPO model with random weights
    let ppo_a = PPO::new(config.clone()).expect("PPO::new should succeed");

    // 2. Save checkpoint
    let tmp = tempfile::tempdir().expect("failed to create tempdir");
    let ckpt_path = tmp.path().join("ppo_checkpoint");

    ppo_a
        .save_checkpoint(&ckpt_path)
        .expect("PPO save_checkpoint should succeed");

    // 3. Load into a fresh PPO model
    let _ppo_b = PPO::load_checkpoint(&ckpt_path).expect("PPO::load_checkpoint should succeed");

    // PPO::load_checkpoint currently re-initializes weights (config-only load).
    // Just verify it doesn't crash and produces valid config.
}

// ---------------------------------------------------------------------------
// PPO: PpoInferenceAdapter deterministic on same weights
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_ppo_inference_adapter_deterministic() {
    let config = small_ppo_config();
    let adapter =
        PpoInferenceAdapter::new(config).expect("PpoInferenceAdapter::new should succeed");

    let fv = fixed_feature_vector_64();
    let pred1 = adapter.predict(&fv).expect("first predict should succeed");
    let pred2 = adapter.predict(&fv).expect("second predict should succeed");

    assert_eq!(
        pred1.direction, pred2.direction,
        "PPO adapter should be deterministic: {} vs {}",
        pred1.direction, pred2.direction,
    );
    assert_eq!(
        pred1.confidence, pred2.confidence,
        "PPO adapter confidence should be deterministic: {} vs {}",
        pred1.confidence, pred2.confidence,
    );
}

// ---------------------------------------------------------------------------
// TFT + Mamba2 Helpers
// ---------------------------------------------------------------------------

/// Small TFT config for fast tests.
/// input_dim = num_static + num_known + num_unknown = 6 + 6 + 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()
    }
}

/// Small Mamba2 config for fast tests.
fn small_mamba2_config() -> Mamba2Config {
    Mamba2Config {
        d_model: 32,
        d_state: 8,
        d_head: 8,
        num_heads: 2,
        expand: 2,
        num_layers: 1,
        max_seq_len: 8,
        dropout: 0.0,
        ..Default::default()
    }
}

const TFT_SEQ_LEN: usize = 4;
const MAMBA2_SEQ_LEN: usize = 4;

/// Build a feature vector with 51 values at a given timestamp.
/// The adapter extracts and zero-pads as needed for its model dimensions.
fn make_fv(ts: i64) -> FeatureVector {
    FeatureVector {
        values: vec![0.1; 51],
        timestamp: ts,
    }
}

/// Feed `count` feature vectors into an adapter to fill its sequence buffer.
/// Returns the predictions from each call (including the neutral ones).
fn fill_adapter_buffer(
    adapter: &dyn ModelInferenceAdapter,
    count: usize,
    base_ts: i64,
) -> Vec<ml::ensemble::inference_adapter::EnsemblePrediction> {
    let mut preds = Vec::with_capacity(count);
    for i in 0..count {
        let fv = make_fv(base_ts + i as i64);
        let pred = adapter.predict(&fv).expect("predict during buffer fill should succeed");
        preds.push(pred);
    }
    preds
}

// ---------------------------------------------------------------------------
// TFT: single adapter deterministic inference (same input => same output)
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_tft_adapter_deterministic_same_buffer() {
    let config = small_tft_config();

    // 1. Create a single TFT adapter
    let adapter = TftInferenceAdapter::new(config, TFT_SEQ_LEN)
        .expect("TftInferenceAdapter creation should succeed");

    // 2. Fill the buffer with identical feature vectors
    let base_ts: i64 = 1_700_000_000;
    for i in 0..TFT_SEQ_LEN {
        let fv = make_fv(base_ts + i as i64);
        let pred = adapter
            .predict(&fv)
            .expect("predict should succeed during fill");

        // While buffering, should return neutral
        if i < TFT_SEQ_LEN - 1 {
            assert_eq!(
                pred.direction, 0.0,
                "Should return neutral direction while buffering (step {i})"
            );
            assert_eq!(
                pred.confidence, 0.0,
                "Should return zero confidence while buffering (step {i})"
            );
        }
    }

    // 3. Adapter should now be ready
    assert!(
        adapter.is_ready(),
        "Adapter should be ready after filling buffer"
    );

    // 4. Feed identical FVs and compare predictions (buffer is full, so each
    //    call runs a real forward pass with the same buffer content)
    let fv = make_fv(base_ts);
    let pred1 = adapter
        .predict(&fv)
        .expect("first full-buffer predict should succeed");
    let pred2 = adapter
        .predict(&fv)
        .expect("second full-buffer predict should succeed");

    // 5. Assert identical outputs (same weights, same buffer content => same output)
    assert!(
        (pred1.direction - pred2.direction).abs() < 1e-6,
        "TFT adapter should be deterministic: direction {} vs {}",
        pred1.direction,
        pred2.direction,
    );
    assert!(
        (pred1.confidence - pred2.confidence).abs() < 1e-6,
        "TFT adapter should be deterministic: confidence {} vs {}",
        pred1.confidence,
        pred2.confidence,
    );

    // 6. Direction and confidence bounds check
    assert!(
        pred1.direction >= -1.0 && pred1.direction <= 1.0,
        "TFT direction {} out of [-1,1]",
        pred1.direction,
    );
    assert!(
        pred1.confidence >= 0.0 && pred1.confidence <= 1.0,
        "TFT confidence {} out of [0,1]",
        pred1.confidence,
    );

    // 7. Model name
    assert_eq!(pred1.model_name, "TFT");
}

// ---------------------------------------------------------------------------
// TFT: quantile metadata verification
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_tft_quantile_metadata_present_after_buffer_fill() {
    let config = small_tft_config();
    let adapter = TftInferenceAdapter::new(config.clone(), TFT_SEQ_LEN)
        .expect("TftInferenceAdapter creation should succeed");

    // Fill the buffer
    let base_ts: i64 = 1_700_000_000;
    let preds = fill_adapter_buffer(&adapter, TFT_SEQ_LEN, base_ts);

    // Predictions before buffer is full should NOT have quantile metadata
    for (i, pred) in preds.iter().enumerate() {
        if i < TFT_SEQ_LEN - 1 {
            assert!(
                pred.metadata.quantiles.is_none(),
                "Neutral prediction at step {i} should not have quantile metadata"
            );
        }
    }

    // The final prediction (buffer just filled) should have quantile metadata
    let last_pred = preds
        .last()
        .expect("should have at least one prediction");
    assert!(
        last_pred.metadata.quantiles.is_some(),
        "TFT prediction with full buffer should include quantile metadata"
    );

    // Verify quantile count matches config.num_quantiles
    let quantiles = last_pred
        .metadata
        .quantiles
        .as_ref()
        .expect("quantiles should be present");
    assert_eq!(
        quantiles.len(),
        config.num_quantiles,
        "Quantile count should match config.num_quantiles ({}), got {}",
        config.num_quantiles,
        quantiles.len(),
    );

    // Quantiles should be finite numbers
    for (i, &q) in quantiles.iter().enumerate() {
        assert!(
            q.is_finite(),
            "Quantile at index {i} should be finite, got {q}"
        );
    }

    // Additional prediction should also have quantiles
    let fv_extra = make_fv(base_ts + TFT_SEQ_LEN as i64);
    let pred_extra = adapter
        .predict(&fv_extra)
        .expect("extra predict should succeed");
    assert!(
        pred_extra.metadata.quantiles.is_some(),
        "Subsequent TFT prediction should also include quantile metadata"
    );
    let quantiles_extra = pred_extra
        .metadata
        .quantiles
        .as_ref()
        .expect("quantiles should be present in extra prediction");
    assert_eq!(
        quantiles_extra.len(),
        config.num_quantiles,
        "Subsequent quantile count should match config"
    );

    // Model name should be "TFT"
    assert_eq!(last_pred.model_name, "TFT");
}

// ---------------------------------------------------------------------------
// Mamba2: single adapter deterministic inference (same input => same output)
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_mamba2_adapter_deterministic_same_buffer() {
    let config = small_mamba2_config();

    // 1. Create a single Mamba2 adapter
    let adapter = Mamba2InferenceAdapter::new(config, MAMBA2_SEQ_LEN)
        .expect("Mamba2InferenceAdapter creation should succeed");

    // 2. Fill the buffer with identical feature vectors
    let base_ts: i64 = 1_700_000_000;
    for i in 0..MAMBA2_SEQ_LEN {
        let fv = make_fv(base_ts);
        let pred = adapter
            .predict(&fv)
            .expect("predict should succeed during fill");

        // While buffering, should return neutral
        if i < MAMBA2_SEQ_LEN - 1 {
            assert_eq!(
                pred.direction, 0.0,
                "Should return neutral direction while buffering (step {i})"
            );
            assert_eq!(
                pred.confidence, 0.0,
                "Should return zero confidence while buffering (step {i})"
            );
        }
    }

    // 3. Adapter should now be ready
    assert!(
        adapter.is_ready(),
        "Adapter should be ready after filling buffer"
    );

    // 4. Feed identical FVs and compare predictions (buffer is full, each
    //    call shifts the ring buffer with the same value, so content is stable)
    let fv = make_fv(base_ts);
    let pred1 = adapter
        .predict(&fv)
        .expect("first full-buffer predict should succeed");
    let pred2 = adapter
        .predict(&fv)
        .expect("second full-buffer predict should succeed");

    // 5. Assert identical outputs (same weights, same buffer content => same output)
    assert!(
        (pred1.direction - pred2.direction).abs() < 1e-6,
        "Mamba2 adapter should be deterministic: direction {} vs {}",
        pred1.direction,
        pred2.direction,
    );
    assert!(
        (pred1.confidence - pred2.confidence).abs() < 1e-6,
        "Mamba2 adapter should be deterministic: confidence {} vs {}",
        pred1.confidence,
        pred2.confidence,
    );

    // 6. Direction and confidence bounds check
    assert!(
        pred1.direction >= -1.0 && pred1.direction <= 1.0,
        "Mamba2 direction {} out of [-1,1]",
        pred1.direction,
    );
    assert!(
        pred1.confidence >= 0.0 && pred1.confidence <= 1.0,
        "Mamba2 confidence {} out of [0,1]",
        pred1.confidence,
    );

    // 7. Model name should be "MAMBA-2"
    assert_eq!(pred1.model_name, "MAMBA-2");

    // 8. Mamba2 should NOT have quantile metadata (that is TFT-specific)
    assert!(
        pred1.metadata.quantiles.is_none(),
        "Mamba2 predictions should not have quantile metadata"
    );
}

// ---------------------------------------------------------------------------
// Mamba2: CheckpointManager + Checkpointable trait roundtrip on Mamba2SSM
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_mamba2_checkpoint_manager_roundtrip() {
    let tmp = tempfile::tempdir().expect("failed to create tempdir");
    let ckpt_config = CheckpointConfig {
        base_dir: tmp.path().to_path_buf(),
        compression: CompressionType::None,
        auto_cleanup: false,
        ..Default::default()
    };

    let manager =
        CheckpointManager::new(ckpt_config).expect("CheckpointManager::new should succeed");

    // Create a Mamba2SSM (implements Checkpointable)
    let config = small_mamba2_config();
    let device = cuda_device();
    let stream = device.cuda_stream().expect("stream").clone();
    let model_a = Mamba2SSM::new(config.clone(), &stream)
        .expect("Mamba2SSM::new should succeed for model_a");

    // Save checkpoint
    let ckpt_id = manager
        .save_checkpoint(&model_a, Some(vec!["mamba2-roundtrip".to_string()]))
        .await
        .expect("save_checkpoint should succeed for Mamba2SSM");

    // Create second model with same config
    let mut model_b =
        Mamba2SSM::new(config, &stream).expect("Mamba2SSM::new should succeed for model_b");

    // Load checkpoint into model_b
    let metadata = manager
        .load_checkpoint(&mut model_b, &ckpt_id)
        .await
        .expect("load_checkpoint should succeed for Mamba2SSM");

    // Verify metadata
    assert!(
        metadata.tags.contains(&"mamba2-roundtrip".to_string()),
        "Checkpoint tags should contain 'mamba2-roundtrip', got {:?}",
        metadata.tags,
    );
    assert_eq!(
        metadata.model_type,
        ml::ModelType::MAMBA,
        "Checkpoint model type should be MAMBA"
    );

    // Verify hyperparameters were preserved
    let d_model_val = metadata
        .hyperparameters
        .get("d_model")
        .and_then(|v| v.as_u64());
    assert_eq!(
        d_model_val,
        Some(32),
        "Checkpoint should preserve d_model=32, got {:?}",
        d_model_val,
    );
}