foxhunt/crates/ml/tests/dqn_inference_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,
)]
//! DQN Checkpoint -> Inference Integration Test
//!
//! Proves the complete checkpoint -> load -> inference path:
//! 1. Train 5 epochs on real market data (fast, just to get a valid checkpoint)
//! 2. Save checkpoint to temp dir via the trainer's checkpoint callback
//! 3. Load checkpoint into a fresh DQN with matching architecture
//! 4. Run inference on 100 synthetic state vectors
//! 5. Assert: actions valid (0..45), Q-values finite, Q-values non-zero

#![allow(unused_crate_dependencies)]

use anyhow::{Context, Result};
// candle eliminated — test uses native cudarc APIs
use ml::dqn::{DQNConfig, DQN};
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use ml_core::cuda_autograd::GpuTensor;
use std::path::PathBuf;
use tracing::info;
use tracing::warn;

/// Locate the small training data directory, returning an error if absent.
fn get_data_dir() -> Result<String> {
    let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .ancestors()
        .find(|p| p.join("test_data").exists())
        .context("Failed to find workspace root with test_data/")?
        .to_path_buf();
    let data_dir = workspace_root.join("test_data/real/databento");
    if !data_dir.exists() {
        anyhow::bail!("Data not found: {}", data_dir.display());
    }
    Ok(data_dir.to_string_lossy().to_string())
}

/// Build a `DQNConfig` that matches the trainer's internal config exactly.
///
/// The trainer builds its DQNConfig from DQNHyperparameters::conservative() with
/// hardcoded architecture params: state_dim=54, num_actions=45, hidden_dims=[256,128,64].
/// NoisyLinear layers are always enabled (unconditional since Rainbow DQN is the
/// only supported architecture). Layer keys use "noisy_hidden_*" / "noisy_output" prefixes.
fn build_matching_config(checkpoint_tensors: &safetensors::SafeTensors<'_>) -> DQNConfig {
    // Discover state_dim from the first NoisyLinear layer weight.
    // NoisyLinear key format: "noisy_hidden_0.mu_w" shape [hidden, input]
    let state_dim = checkpoint_tensors
        .tensors()
        .iter()
        .find(|(name, _)| name.contains("noisy_hidden_0") && name.contains("mu_w"))
        .map(|(_, view)| {
            let d = view.shape();
            if d.len() == 2 { d[1] } else { 54 }
        })
        .unwrap_or(54);

    // Replicate the exact config the trainer constructs (see trainer.rs ~line 289).
    let _ = state_dim; // state_dim is now the STATE_DIM constant (no longer configurable)
    let mut config = DQNConfig::conservative();
    config.num_actions = 45;
    config.hidden_dims = vec![256, 128, 64];
    // Trainer hardcodes noisy_sigma_init = 0.5 via hyperparams
    config.noisy_sigma_init = 0.5;
    // Match trainer defaults: CQL enabled with alpha=0.1
    config.cql_alpha = 0.1;
    config
}

#[tokio::test]
async fn test_checkpoint_to_inference() -> Result<()> {
    // -- Skip gracefully when real data is absent (CI environments) --------
    let data_dir = match get_data_dir() {
        Ok(dir) => dir,
        Err(e) => {
            warn!(reason = %e, "Skipping test: data not available");
            return Ok(());
        }
    };

    let checkpoint_dir = tempfile::tempdir()?;

    // =====================================================================
    // Phase 1: Train for 5 epochs to produce a valid checkpoint
    // =====================================================================
    let mut hyperparams = DQNHyperparameters::conservative();
    hyperparams.replay_buffer_vram_fraction = 0.0; // Disable AutoReplaySizer for test determinism
    hyperparams.epochs = 5;
    hyperparams.batch_size = 32;
    hyperparams.learning_rate = 0.0001;
    hyperparams.early_stopping_enabled = false;
    hyperparams.checkpoint_frequency = 5; // checkpoint on last epoch

    let mut trainer = DQNTrainer::new(hyperparams)?;
    // The closure captures shared state across the async-checkpoint worker
    // boundary, so wrap in Arc<Mutex<>> to satisfy the `+ 'static` bound
    // (worker can't borrow the test fn's stack frame).
    let best_checkpoint_path: std::sync::Arc<std::sync::Mutex<Option<PathBuf>>> =
        std::sync::Arc::new(std::sync::Mutex::new(None));
    let best_cb_handle = std::sync::Arc::clone(&best_checkpoint_path);
    let ckpt_dir_path = checkpoint_dir.path().to_path_buf();

    let _metrics = trainer
        .train(&data_dir, "ES.FUT", move |epoch, checkpoint_data, is_best| {
            let name = if is_best {
                "inference_best.safetensors".to_string()
            } else {
                format!("inference_epoch_{epoch}.safetensors")
            };
            let path = ckpt_dir_path.join(&name);
            std::fs::write(&path, &checkpoint_data)?;
            if is_best {
                *best_cb_handle.lock().unwrap() = Some(path.clone());
            }
            Ok(path.to_string_lossy().to_string())
        })
        .await?;

    // If no "best" was saved, fall back to any checkpoint in the directory.
    let best_checkpoint_path = best_checkpoint_path.lock().unwrap().clone();
    let checkpoint_path = match best_checkpoint_path {
        Some(p) => p,
        None => {
            // Find first .safetensors file in the temp dir
            let mut entries: Vec<_> = std::fs::read_dir(checkpoint_dir.path())?
                .filter_map(|e| e.ok())
                .filter(|e| {
                    e.path()
                        .extension()
                        .map_or(false, |ext| ext == "safetensors")
                })
                .collect();
            anyhow::ensure!(
                !entries.is_empty(),
                "No checkpoint files were saved during training"
            );
            entries.sort_by_key(|e| e.path());
            entries.pop().context("No checkpoint file")?.path()
        }
    };

    assert!(
        checkpoint_path.exists(),
        "Checkpoint file does not exist: {}",
        checkpoint_path.display()
    );
    info!(
        checkpoint = %checkpoint_path.display(),
        bytes = std::fs::metadata(&checkpoint_path)?.len(),
        "Phase 1 complete: checkpoint saved"
    );

    // =====================================================================
    // Phase 2: Load checkpoint into a fresh DQN
    // =====================================================================
    // Load the checkpoint's tensor metadata to discover architecture params
    // (state_dim, noisy vs standard layers) so we can construct a DQN whose
    // VarMap key names and tensor shapes match exactly.
    let checkpoint_path_str = checkpoint_path
        .to_str()
        .context("Non-UTF8 checkpoint path")?;

    let checkpoint_bytes = std::fs::read(checkpoint_path_str)?;
    let checkpoint_tensors = safetensors::SafeTensors::deserialize(&checkpoint_bytes)?;

    info!(
        tensor_count = checkpoint_tensors.len(),
        keys = ?checkpoint_tensors.names(),
        "Checkpoint tensor inventory"
    );

    let config = build_matching_config(&checkpoint_tensors);
    let state_dim = ml_core::state_layout::STATE_DIM;
    info!(
        state_dim = ml_core::state_layout::STATE_DIM,
        num_actions = config.num_actions,
        "Built matching config (noisy nets always enabled)"
    );

    let mut fresh_dqn = DQN::new(config)?;
    fresh_dqn.load_from_safetensors(checkpoint_path_str)?;

    info!("Phase 2 complete: fresh DQN loaded from checkpoint");

    // =====================================================================
    // Phase 3: Inference on 100 synthetic state vectors
    // =====================================================================
    let num_inference_samples: usize = 100;
    let num_actions: usize = 45;
    let stream = fresh_dqn.cuda_stream();

    let mut all_actions_valid = true;
    let mut all_q_finite = true;
    let mut any_q_nonzero = false;
    let mut action_counts = vec![0usize; num_actions];

    for i in 0..num_inference_samples {
        // Deterministic synthetic state: slight variation per sample
        let state_vec: Vec<f32> = (0..state_dim)
            .map(|j| ((i * state_dim + j) as f32 * 0.01).sin())
            .collect();

        let state_tensor = GpuTensor::from_host(&state_vec, vec![1, state_dim], stream)?;

        let q_values = fresh_dqn.forward(&state_tensor)?;

        // Shape check: [1, 45]
        let dims = q_values.shape();
        assert_eq!(
            dims,
            &[1, num_actions],
            "Q-value tensor shape mismatch: expected [1, {}], got {:?}",
            num_actions,
            dims
        );

        let q_vec: Vec<f32> = q_values.to_host(stream)?;

        // Check finite
        for &q in &q_vec {
            if !q.is_finite() {
                all_q_finite = false;
            }
            if q.abs() > 1e-9 {
                any_q_nonzero = true;
            }
        }

        // Argmax to get action index
        let best_action_idx = q_vec
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(idx, _)| idx)
            .unwrap_or(0);

        if best_action_idx >= num_actions {
            all_actions_valid = false;
        } else if let Some(count) = action_counts.get_mut(best_action_idx) {
            *count += 1;
        }
    }

    // =====================================================================
    // Phase 4: Assertions
    // =====================================================================
    assert!(all_q_finite, "Some Q-values were NaN or Inf");
    assert!(
        any_q_nonzero,
        "All Q-values were zero -- model likely failed to load weights"
    );
    assert!(
        all_actions_valid,
        "Some action indices were out of range (expected 0..45)"
    );

    // Report action distribution
    let unique_actions = action_counts.iter().filter(|&&c| c > 0).count();
    info!(
        inference_passes = num_inference_samples,
        unique_actions,
        "Phase 3 complete: inference passes done"
    );

    // Sanity: at least 1 unique action (degenerate model is still valid for this test)
    assert!(
        unique_actions >= 1,
        "No actions were selected -- inference loop failure"
    );

    info!("All assertions passed: checkpoint -> load -> inference pipeline verified");

    Ok(())
}