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::cuda_autograd::GpuTensor;
use ml::dqn::agent::DQNAgent;
use ml::dqn::dqn::{DQNConfig, DQN};
use ml::ensemble::adapters::{
    DqnInferenceAdapter, 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_dqn_config() -> DQNConfig {
    DQNConfig {
        state_dim: 51,
        num_actions: 45,
        hidden_dims: vec![32, 32],
        use_iqn: false,
        ..Default::default()
    }
}

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_51() -> FeatureVector {
    FeatureVector {
        values: vec![0.3; 51],
        timestamp: 1_700_000_000,
    }
}

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

/// Extract Q-values from a DQN model given an f32 input vector.
///
/// Returns the raw Q-value vector (f64) for comparison.
fn dqn_q_values(model: &DQN, input: &[f32]) -> Result<Vec<f64>, ml::MLError> {
    let device = cuda_device();
    let stream = device.cuda_stream()?.clone();
    let len = input.len();
    let tensor = GpuTensor::from_host(input, vec![1, len], &stream)
        .map_err(|e| ml::MLError::ModelError(format!("Failed to create tensor: {e}")))?;
    let output = model.forward(&tensor)?;
    let squeezed = output
        .squeeze(0, &stream)
        .map_err(|e| ml::MLError::ModelError(format!("Failed to squeeze: {e}")))?;
    let f32_vec: Vec<f32> = squeezed
        .to_vec1(&stream)
        .map_err(|e| ml::MLError::ModelError(format!("Failed to extract vec: {e}")))?;
    Ok(f32_vec.iter().map(|&v| v as f64).collect())
}

/// Save GpuVarStore weights to a safetensors file on disk.
fn save_varstore_safetensors(
    vars: &ml::cuda_autograd::GpuVarStore,
    path: &std::path::Path,
) {
    let host_data = vars.export_to_host().expect("export_to_host should succeed");

    // Build owned byte buffers for each tensor
    let owned: Vec<(String, Vec<usize>, Vec<u8>)> = host_data
        .iter()
        .map(|(name, (shape, values))| {
            let bytes: Vec<u8> = values.iter().flat_map(|v| v.to_le_bytes()).collect();
            (name.clone(), shape.clone(), bytes)
        })
        .collect();

    // Build TensorView references
    let views: Vec<(&str, safetensors::tensor::TensorView<'_>)> = owned
        .iter()
        .map(|(name, shape, bytes)| {
            (
                name.as_str(),
                safetensors::tensor::TensorView::new(
                    safetensors::Dtype::F32,
                    shape.clone(),
                    bytes,
                )
                .expect("TensorView"),
            )
        })
        .collect();

    let serialized =
        safetensors::tensor::serialize(views, None).expect("serialize should succeed");
    std::fs::write(path, serialized).expect("write should succeed");
}

// ---------------------------------------------------------------------------
// DQN: raw Q-network weight roundtrip
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_dqn_raw_weight_roundtrip_exact_output() {
    let config = small_dqn_config();

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

    // 2. Run inference to get Q-values before save
    let input = vec![0.3f32; 51];
    let q_before = dqn_q_values(&model_a, &input).expect("forward should succeed");
    assert!(!q_before.is_empty(), "Q-values should not be empty");

    // 3. Save weights to temp directory via export_to_host + safetensors
    let tmp = tempfile::tempdir().expect("failed to create tempdir");
    let ckpt_path = tmp.path().join("dqn_roundtrip.safetensors");
    save_varstore_safetensors(model_a.get_q_network_vars(), &ckpt_path);

    // 4. Create second DQN model (same config, different random init)
    let mut model_b = DQN::new(config).expect("DQN::new should succeed for model_b");

    // 5. Load weights from checkpoint
    let path_str = ckpt_path
        .to_str()
        .expect("temp path should be valid UTF-8");
    model_b
        .load_from_safetensors(path_str)
        .expect("load_from_safetensors should succeed");

    // 6. Run inference on same input
    let q_after = dqn_q_values(&model_b, &input).expect("forward after load should succeed");

    // 7. Assert exact element-wise equality
    assert_eq!(
        q_before.len(),
        q_after.len(),
        "Q-value vector length mismatch: {} vs {}",
        q_before.len(),
        q_after.len(),
    );

    for (i, (a, b)) in q_before.iter().zip(q_after.iter()).enumerate() {
        assert!(
            (a - b).abs() < 1e-6,
            "Q-value mismatch at index {i}: {a} vs {b}",
        );
    }
}

// ---------------------------------------------------------------------------
// DQN: adapter roundtrip via from_checkpoint
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_dqn_adapter_roundtrip_via_from_checkpoint() {
    let config = small_dqn_config();

    // 1. Create a DQN model, save its weights via export_to_host + safetensors
    let model = DQN::new(config.clone()).expect("DQN::new should succeed");
    let tmp = tempfile::tempdir().expect("failed to create tempdir");
    let ckpt_path = tmp.path().join("dqn_adapter_test.safetensors");
    save_varstore_safetensors(model.get_q_network_vars(), &ckpt_path);

    // 2. Get reference Q-values from original model
    let input = vec![0.3f32; 51];
    let q_original = dqn_q_values(&model, &input).expect("original forward should succeed");

    // 3. Load into a DqnInferenceAdapter via from_checkpoint
    let path_str = ckpt_path
        .to_str()
        .expect("temp path should be valid UTF-8");
    let adapter = DqnInferenceAdapter::from_checkpoint(config, path_str)
        .expect("from_checkpoint should succeed");

    // 4. Run inference through the adapter
    let fv = fixed_feature_vector_51();
    let pred = adapter.predict(&fv).expect("adapter predict should succeed");

    // 5. Verify prediction shape and ranges
    assert!(
        pred.direction >= -1.0 && pred.direction <= 1.0,
        "direction {} out of [-1,1]",
        pred.direction,
    );
    assert!(
        pred.confidence >= 0.0 && pred.confidence <= 1.0,
        "confidence {} out of [0,1]",
        pred.confidence,
    );

    // 6. Verify Q-values from adapter match original model
    let adapter_q = pred
        .metadata
        .q_values
        .as_ref()
        .expect("DQN adapter should emit q_values");
    assert_eq!(
        q_original.len(),
        adapter_q.len(),
        "Q-value vector length mismatch",
    );
    for (i, (orig, loaded)) in q_original.iter().zip(adapter_q.iter()).enumerate() {
        assert!(
            (orig - loaded).abs() < 1e-5,
            "Q-value mismatch at index {i}: original={orig} vs loaded={loaded}",
        );
    }
}

// ---------------------------------------------------------------------------
// DQN: CheckpointManager + Checkpointable trait roundtrip
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_dqn_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 DQNAgent (the type that implements Checkpointable)
    let config = small_dqn_config();
    let agent_a = DQNAgent::new(config.clone()).expect("DQNAgent::new should succeed");

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

    // Create second agent with same config
    let mut agent_b = DQNAgent::new(config).expect("DQNAgent::new should succeed for agent_b");

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

    assert_eq!(metadata.model_name, "dqn_agent");
    assert!(metadata.tags.contains(&"roundtrip-test".to_string()));
}

// ---------------------------------------------------------------------------
// 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,
    );
}