foxhunt/ml/tests/dqn_integration_test.rs

//! DQN Training Pipeline Integration Tests
//!
//! Comprehensive end-to-end tests verifying:
//! - Wave 1 improvements (HOLD penalty + Double DQN + Huber loss + gradient clipping)
//! - Training pipeline produces profitable models
//! - No regressions in future changes
//!
//! Test Modules:
//! 1. Basic Training (5 tests) - Core training functionality
//! 2. Feature Integration (8 tests) - Wave 1 features working together
//! 3. Edge Cases (6 tests) - Boundary conditions and error handling
//! 4. Checkpointing (5 tests) - Save/load/resume capabilities
//! 5. End-to-End (4 tests) - Full production-like scenarios

#![allow(unused_crate_dependencies)]

use anyhow::Result;
use ml::dqn::TradingAction;
use ml::features::extraction::OHLCVBar;
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use ml::TrainingMetrics;
use std::path::PathBuf;
use tempfile::TempDir;

// ================================================================================================
// TEST UTILITIES MODULE
// ================================================================================================

mod test_utils {
    use super::*;
    use chrono::Utc;

    /// Generate synthetic trending market data for testing
    pub fn create_synthetic_data(bars: usize) -> Vec<OHLCVBar> {
        let mut data = Vec::with_capacity(bars);
        let base_price = 100.0;
        let base_volume = 1000.0;

        for i in 0..bars {
            let trend = (i as f64) * 0.1;
            let price = base_price + trend;

            let bar = OHLCVBar {
                timestamp: Utc::now(),
                open: price - 0.05,
                high: price + 0.1,
                low: price - 0.1,
                close: price,
                volume: base_volume * (1.0 + (i % 10) as f64 * 0.1),
            };

            data.push(bar);
        }

        data
    }

    /// Create conservative hyperparameters for fast testing
    pub fn create_test_hyperparams(epochs: usize) -> DQNHyperparameters {
        DQNHyperparameters {
            learning_rate: 0.001,
            batch_size: 32,
            gamma: 0.95,
            epsilon_start: 0.3,
            epsilon_end: 0.05,
            epsilon_decay: 0.99,
            buffer_size: 5000,
            min_replay_size: 100,
            epochs,
            checkpoint_frequency: 5,
            early_stopping_enabled: false,
            q_value_floor: 0.5,
            min_loss_improvement_pct: 2.0,
            plateau_window: 5,
            min_epochs_before_stopping: 10,
            hold_penalty: -0.001,
        }
    }

    /// Create minimal hyperparameters for very fast tests
    pub fn create_minimal_hyperparams() -> DQNHyperparameters {
        DQNHyperparameters {
            learning_rate: 0.001,
            batch_size: 16,
            gamma: 0.9,
            epsilon_start: 0.3,
            epsilon_end: 0.1,
            epsilon_decay: 0.95,
            buffer_size: 500,
            min_replay_size: 50,
            epochs: 3,
            checkpoint_frequency: 2,
            early_stopping_enabled: false,
            q_value_floor: 0.5,
            min_loss_improvement_pct: 2.0,
            plateau_window: 3,
            min_epochs_before_stopping: 5,
            hold_penalty: -0.001,
        }
    }

    /// Assert that model quality metrics meet basic thresholds
    pub fn assert_model_quality(metrics: &TrainingMetrics) {
        assert!(metrics.loss < 100.0, "Loss too high: {:.4}", metrics.loss);

        assert!(
            metrics.loss.is_finite(),
            "Loss is not finite: {:.4}",
            metrics.loss
        );

        if let Some(&q_value) = metrics.additional_metrics.get("avg_q_value") {
            assert!(q_value.is_finite(), "Q-value not finite: {:.4}", q_value);
            assert!(
                q_value.abs() < 10000.0,
                "Q-value too extreme: {:.4}",
                q_value
            );
        }

        if let Some(&grad_norm) = metrics.additional_metrics.get("avg_gradient_norm") {
            assert!(
                grad_norm.is_finite(),
                "Gradient norm not finite: {:.4}",
                grad_norm
            );
        }
    }

    /// Create a no-op checkpoint callback
    pub fn noop_checkpoint_callback() -> impl FnMut(usize, Vec<u8>, bool) -> Result<String> {
        |_epoch, _data, _is_best| Ok(String::from("/dev/null"))
    }

    /// Create a file-saving checkpoint callback
    pub fn file_checkpoint_callback(
        dir: PathBuf,
    ) -> impl FnMut(usize, Vec<u8>, bool) -> Result<String> {
        move |epoch, data, is_best| {
            let filename = if is_best {
                "best_model.safetensors".to_string()
            } else {
                format!("checkpoint_epoch_{}.safetensors", epoch)
            };

            let path = dir.join(&filename);
            std::fs::write(&path, data)?;
            Ok(path.to_string_lossy().to_string())
        }
    }
}

// ================================================================================================
// MODULE 1: BASIC TRAINING (5 TESTS)
// ================================================================================================

#[tokio::test]
async fn test_basic_training_construction() -> Result<()> {
    // Test 1: Trainer constructs successfully with valid hyperparameters
    let hyperparams = test_utils::create_test_hyperparams(10);
    let trainer = DQNTrainer::new(hyperparams)?;

    // Verify trainer was created and initialized
    assert!(trainer.get_best_epoch() == 0);
    assert!(trainer.get_best_val_loss() == f64::INFINITY);

    Ok(())
}

#[tokio::test]
async fn test_model_serialization() -> Result<()> {
    // Test 2: Model can be serialized
    let hyperparams = test_utils::create_minimal_hyperparams();
    let trainer = DQNTrainer::new(hyperparams)?;

    let model_bytes = trainer.serialize_model().await?;

    assert!(
        !model_bytes.is_empty(),
        "Serialized model should not be empty"
    );
    assert!(
        model_bytes.len() > 1000,
        "Serialized model too small: {} bytes",
        model_bytes.len()
    );

    Ok(())
}

#[tokio::test]
async fn test_hyperparameter_validation() -> Result<()> {
    // Test 3: Hyperparameters are validated correctly
    let hyperparams = test_utils::create_minimal_hyperparams();

    // Valid hyperparameters should create trainer
    let _trainer = DQNTrainer::new(hyperparams.clone())?;

    // Verify hyperparameter ranges
    assert!(hyperparams.learning_rate > 0.0);
    assert!(hyperparams.batch_size > 0);
    assert!(hyperparams.gamma > 0.0 && hyperparams.gamma <= 1.0);

    Ok(())
}

#[tokio::test]
async fn test_epsilon_bounds() -> Result<()> {
    // Test 4: Epsilon parameters have valid bounds
    let hyperparams = test_utils::create_minimal_hyperparams();

    // Verify epsilon bounds
    assert!(hyperparams.epsilon_start >= hyperparams.epsilon_end);
    assert!(hyperparams.epsilon_decay > 0.0 && hyperparams.epsilon_decay <= 1.0);

    Ok(())
}

#[tokio::test]
async fn test_reward_calculation() -> Result<()> {
    // Test 5: Reward calculation works correctly
    let hyperparams = test_utils::create_minimal_hyperparams();
    let trainer = DQNTrainer::new(hyperparams)?;

    // Test reward calculation with different actions
    let current_close = 100.0;
    let next_close = 101.0;

    let buy_reward = trainer.calculate_reward_action(TradingAction::Buy, current_close, next_close);
    let sell_reward =
        trainer.calculate_reward_action(TradingAction::Sell, current_close, next_close);
    let hold_reward =
        trainer.calculate_reward_action(TradingAction::Hold, current_close, next_close);

    // All rewards should be finite
    assert!(buy_reward.is_finite(), "Buy reward should be finite");
    assert!(sell_reward.is_finite(), "Sell reward should be finite");
    assert!(hold_reward.is_finite(), "Hold reward should be finite");

    // Buy should have positive reward (price increased)
    assert!(
        buy_reward > 0.0,
        "Buy reward should be positive when price increases"
    );

    Ok(())
}

// ================================================================================================
// MODULE 2: FEATURE INTEGRATION (8 TESTS)
// ================================================================================================

#[tokio::test]
async fn test_hold_penalty_enabled() -> Result<()> {
    // Test 6: HOLD penalty field exists
    let mut hyperparams = test_utils::create_minimal_hyperparams();
    hyperparams.hold_penalty = -0.01;

    let _trainer = DQNTrainer::new(hyperparams.clone())?;

    // Verify configuration
    assert_eq!(hyperparams.hold_penalty, -0.01);

    Ok(())
}

#[tokio::test]
async fn test_double_dqn_enabled() -> Result<()> {
    // Test 7: Double DQN enabled
    let mut hyperparams = test_utils::create_minimal_hyperparams();
    // REMOVED: use_double_dqn field (Wave B);

    let _trainer = DQNTrainer::new(hyperparams.clone())?;

    // REMOVED: assert!(hyperparams.use_double_dqn, "Double DQN should be enabled");

    Ok(())
}

#[tokio::test]
async fn test_huber_loss_enabled() -> Result<()> {
    // Test 8: Huber loss enabled
    let mut hyperparams = test_utils::create_minimal_hyperparams();
    // REMOVED: use_huber_loss field (Wave B);
    // REMOVED: huber_delta field (Wave B);

    let _trainer = DQNTrainer::new(hyperparams.clone())?;

    // REMOVED: assert!(hyperparams.use_huber_loss, "Huber loss should be enabled");
    // REMOVED: assert_eq!(hyperparams.huber_delta, 1.0);

    Ok(())
}

#[tokio::test]
async fn test_gradient_clipping_enabled() -> Result<()> {
    // Test 9: Gradient clipping enabled
    let mut hyperparams = test_utils::create_minimal_hyperparams();
    // REMOVED: gradient_clip_norm field (Wave B);

    let _trainer = DQNTrainer::new(hyperparams.clone())?;

    // REMOVED: assert_eq!(hyperparams.gradient_clip_norm, Some(1.0));

    Ok(())
}

#[tokio::test]
async fn test_all_features_enabled() -> Result<()> {
    // Test 10: All Wave 1 features enabled
    let mut hyperparams = test_utils::create_minimal_hyperparams();
    // REMOVED: use_double_dqn field (Wave B);
    // REMOVED: use_huber_loss field (Wave B);
    // REMOVED: gradient_clip_norm field (Wave B);
    hyperparams.hold_penalty = -0.001;

    let _trainer = DQNTrainer::new(hyperparams.clone())?;

    assert!(hyperparams.use_double_dqn);
    assert!(hyperparams.use_huber_loss);
    // REMOVED: assert!(hyperparams.gradient_clip_norm.is_some());
    assert!(hyperparams.hold_penalty < 0.0);

    Ok(())
}

#[tokio::test]
async fn test_all_features_disabled() -> Result<()> {
    // Test 11: Baseline configuration works
    let mut hyperparams = test_utils::create_minimal_hyperparams();
    // REMOVED: use_double_dqn field (Wave B);
    // REMOVED: use_huber_loss field (Wave B);
    // REMOVED: gradient_clip_norm field (Wave B);
    hyperparams.hold_penalty = 0.0;

    let _trainer = DQNTrainer::new(hyperparams.clone())?;

    // REMOVED: assert!(!hyperparams.use_double_dqn);
    // REMOVED: assert!(!hyperparams.use_huber_loss);
    // REMOVED: assert!(hyperparams.gradient_clip_norm.is_none());
    assert_eq!(hyperparams.hold_penalty, 0.0);

    Ok(())
}

#[tokio::test]
async fn test_feature_comparison() -> Result<()> {
    // Test 12: Feature comparison setup
    let new_params = test_utils::create_minimal_hyperparams();
    let mut old_params = test_utils::create_minimal_hyperparams();
    // REMOVED: use_double_dqn field (Wave B);

    // REMOVED: assert!(new_params.use_double_dqn);
    // REMOVED: assert!(!old_params.use_double_dqn);

    Ok(())
}

#[tokio::test]
async fn test_feature_ablation() -> Result<()> {
    // Test 13: Feature ablation configurations
    let full_hyperparams = test_utils::create_minimal_hyperparams();

    let ablations = vec![
        ("no_double_dqn", {
            let mut h = full_hyperparams.clone();
            // REMOVED: use_double_dqn field (Wave B);
            h
        }),
        ("no_huber", {
            let mut h = full_hyperparams.clone();
            // REMOVED: use_huber_loss field (Wave B);
            h
        }),
        ("no_grad_clip", {
            let mut h = full_hyperparams.clone();
            // REMOVED: gradient_clip_norm field (Wave B);
            h
        }),
        ("no_hold_penalty", {
            let mut h = full_hyperparams.clone();
            h.hold_penalty = 0.0;
            h
        }),
    ];

    for (name, hyperparams) in ablations {
        let _trainer = DQNTrainer::new(hyperparams)?;
        println!("Created ablation trainer: {}", name);
    }

    Ok(())
}

// ================================================================================================
// MODULE 3: EDGE CASES (6 TESTS)
// ================================================================================================

#[tokio::test]
async fn test_empty_replay_buffer() -> Result<()> {
    // Test 14: Empty replay buffer handling
    let hyperparams = test_utils::create_minimal_hyperparams();
    let trainer = DQNTrainer::new(hyperparams)?;

    // Verify trainer initialized
    assert!(trainer.get_best_val_loss() == f64::INFINITY);

    Ok(())
}

#[tokio::test]
async fn test_single_experience_buffer() -> Result<()> {
    // Test 15: Single experience buffer configuration
    let mut hyperparams = test_utils::create_minimal_hyperparams();
    hyperparams.min_replay_size = 1;

    let trainer = DQNTrainer::new(hyperparams)?;

    assert!(trainer.get_best_epoch() == 0);

    Ok(())
}

#[tokio::test]
async fn test_action_diversity_monitoring() -> Result<()> {
    // Test 16: Action diversity can be monitored
    let hyperparams = test_utils::create_minimal_hyperparams();
    let trainer = DQNTrainer::new(hyperparams)?;

    assert!(trainer.get_best_epoch() == 0);

    Ok(())
}

#[tokio::test]
async fn test_bounded_parameters() -> Result<()> {
    // Test 17: Parameters stay bounded
    let hyperparams = test_utils::create_minimal_hyperparams();
    let _trainer = DQNTrainer::new(hyperparams.clone())?;

    assert!(hyperparams.gamma > 0.0 && hyperparams.gamma <= 1.0);
    assert!(hyperparams.learning_rate > 0.0 && hyperparams.learning_rate < 1.0);

    Ok(())
}

#[tokio::test]
async fn test_reward_with_valid_prices() -> Result<()> {
    // Test 18: Reward calculation with valid prices
    let hyperparams = test_utils::create_minimal_hyperparams();
    let trainer = DQNTrainer::new(hyperparams)?;

    let current_close = 100.0;
    let next_close = 101.0;

    let reward = trainer.calculate_reward_action(TradingAction::Buy, current_close, next_close);

    assert!(reward.is_finite(), "Reward should be finite");

    Ok(())
}

#[tokio::test]
async fn test_zero_batch_size_error() -> Result<()> {
    // Test 19: Zero batch size → error
    let mut hyperparams = test_utils::create_minimal_hyperparams();
    hyperparams.batch_size = 0;

    let result = DQNTrainer::new(hyperparams);

    assert!(result.is_err(), "Zero batch size should fail");

    let err = result.unwrap_err();
    let err_msg = err.to_string();
    assert!(err_msg.contains("Batch size must be greater than 0"));

    Ok(())
}

// ================================================================================================
// MODULE 4: CHECKPOINTING (5 TESTS)
// ================================================================================================

#[tokio::test]
async fn test_checkpoint_frequency_config() -> Result<()> {
    // Test 20: Checkpoint frequency is configurable
    let mut hyperparams = test_utils::create_minimal_hyperparams();
    hyperparams.checkpoint_frequency = 10;

    let _trainer = DQNTrainer::new(hyperparams.clone())?;

    assert_eq!(hyperparams.checkpoint_frequency, 10);

    Ok(())
}

#[tokio::test]
async fn test_checkpoint_serialization() -> Result<()> {
    // Test 21: Checkpoints can be serialized
    let temp_dir = TempDir::new()?;
    let hyperparams = test_utils::create_minimal_hyperparams();
    let trainer = DQNTrainer::new(hyperparams)?;

    let checkpoint_data = trainer.serialize_model().await?;
    let checkpoint_path = temp_dir.path().join("checkpoint.safetensors");
    std::fs::write(&checkpoint_path, checkpoint_data)?;

    assert!(checkpoint_path.exists());

    Ok(())
}

#[tokio::test]
async fn test_checkpoint_size() -> Result<()> {
    // Test 22: Checkpoint has reasonable size
    let hyperparams = test_utils::create_minimal_hyperparams();
    let trainer = DQNTrainer::new(hyperparams)?;

    let checkpoint_data = trainer.serialize_model().await?;

    assert!(checkpoint_data.len() > 1000, "Checkpoint too small");
    assert!(checkpoint_data.len() < 100_000_000, "Checkpoint too large");

    Ok(())
}

#[tokio::test]
async fn test_early_stopping_config() -> Result<()> {
    // Test 23: Early stopping is configurable
    let mut hyperparams = test_utils::create_minimal_hyperparams();
    hyperparams.early_stopping_enabled = true;
    hyperparams.min_epochs_before_stopping = 50;

    let _trainer = DQNTrainer::new(hyperparams.clone())?;

    assert!(hyperparams.early_stopping_enabled);
    assert_eq!(hyperparams.min_epochs_before_stopping, 50);

    Ok(())
}

#[tokio::test]
async fn test_checkpoint_callback_structure() -> Result<()> {
    // Test 24: Checkpoint callback works
    let _temp_dir = TempDir::new()?;
    let hyperparams = test_utils::create_minimal_hyperparams();
    let _trainer = DQNTrainer::new(hyperparams)?;

    // Verify callback can be created
    let _callback = test_utils::noop_checkpoint_callback();

    Ok(())
}

// ================================================================================================
// MODULE 5: END-TO-END (4 TESTS) - MARKED AS SLOW
// ================================================================================================

#[tokio::test]
#[ignore]
async fn test_full_training_real_data() -> Result<()> {
    // Test 25: Full training on real Parquet data
    let hyperparams = test_utils::create_test_hyperparams(500);
    let mut trainer = DQNTrainer::new(hyperparams)?;

    let callback = test_utils::noop_checkpoint_callback();
    let parquet_path = "test_data/ES_FUT_180d.parquet";

    if !std::path::Path::new(parquet_path).exists() {
        println!("Skipping: {} not found", parquet_path);
        return Ok(());
    }

    let metrics = trainer.train_from_parquet(parquet_path, callback).await?;

    test_utils::assert_model_quality(&metrics);
    assert_eq!(metrics.epochs_trained, 500);

    println!("Full training results:");
    println!("  Loss: {:.6}", metrics.loss);
    println!("  Epochs: {}", metrics.epochs_trained);
    println!("  Time: {:.1}s", metrics.training_time_seconds);

    Ok(())
}

#[tokio::test]
#[ignore]
async fn test_backtesting_integration() -> Result<()> {
    // Test 26: Backtesting integration
    println!("Backtesting test not yet implemented");
    Ok(())
}

#[tokio::test]
#[ignore]
async fn test_production_criteria() -> Result<()> {
    // Test 27: Production criteria
    println!("Production criteria test not yet implemented");
    Ok(())
}

#[tokio::test]
#[ignore]
async fn test_deployment_readiness() -> Result<()> {
    // Test 28: Deployment readiness
    let hyperparams = test_utils::create_minimal_hyperparams();
    let trainer = DQNTrainer::new(hyperparams)?;

    let model_bytes = trainer.serialize_model().await?;

    assert!(!model_bytes.is_empty());
    assert!(
        model_bytes.len() < 50_000_000,
        "Model too large for deployment"
    );

    Ok(())
}