foxhunt/crates/ml/tests/dqn_episode_boundaries_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,
)]
//! P1 Episode Boundaries Fix - Comprehensive Test Suite
//!
//! This test suite validates the episode boundary segmentation fix that addresses
//! the missing episode termination logic causing improper value bootstrapping.
//!
//! **Problem**: Entire 18,000-sample dataset treated as SINGLE episode
//! **Fix**: 200-sample episodes with portfolio resets at boundaries
//! **Impact**: 1 → 90 terminal states per epoch (90x improvement)

use ml::dqn::portfolio_tracker::PortfolioTracker;
use ml::dqn::agent::TradingAction;
use tracing::info;

/// Episode length constant (must match dqn.rs)
const EPISODE_LENGTH: usize = 200;

#[test]
fn test_episode_boundary_detection() {
    // Test 1: Verify episode boundaries occur at correct sample indices
    //
    // For 18,000 samples with EPISODE_LENGTH=200:
    // Expected boundaries: 200, 400, 600, ..., 17800, 18000
    // Total boundaries: 90
    
    let dataset_size = 18_000;
    let expected_boundaries = 90;
    
    let mut boundary_count = 0;
    let mut boundary_indices = Vec::new();
    
    for i in 0..dataset_size {
        let done = (i + 1) % EPISODE_LENGTH == 0 || i + 1 >= dataset_size;
        if done {
            boundary_count += 1;
            boundary_indices.push(i + 1);
        }
    }
    
    assert_eq!(boundary_count, expected_boundaries, 
        "Expected {} episode boundaries, got {}", expected_boundaries, boundary_count);
    
    // Verify first few boundaries
    assert_eq!(boundary_indices[0], 200, "First boundary should be at sample 200");
    assert_eq!(boundary_indices[1], 400, "Second boundary should be at sample 400");
    assert_eq!(boundary_indices[2], 600, "Third boundary should be at sample 600");
    
    // Verify last boundary
    assert_eq!(boundary_indices[boundary_indices.len() - 1], 18_000, 
        "Last boundary should be at final sample");
        
    info!(boundary_count, "Test 1 PASSED: Episode boundaries detected correctly");
}

#[test]
fn test_portfolio_reset_at_boundaries() {
    // Test 2: Verify portfolio state resets at episode boundaries
    //
    // Portfolio should reset to:
    // - position_size = 0.0 (flat)
    // - cash = initial_capital
    // - entry_price = 0.0
    
    let mut tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);

    // Simulate trading activity (open position)
    tracker.execute_legacy_action(TradingAction::Buy, 100.0, 10.0);
    
    // Verify position is open
    assert_eq!(tracker.current_position(), 10.0, "Position should be 10.0 before reset");
    assert_eq!(tracker.cash_balance(), 9_000.0, "Cash should be 9000.0 before reset");
    
    // Reset at episode boundary
    tracker.reset();
    
    // Verify portfolio reset correctly
    assert_eq!(tracker.current_position(), 0.0, "Position should be 0.0 after reset");
    assert_eq!(tracker.cash_balance(), 10_000.0, "Cash should be 10000.0 after reset");
    assert_eq!(tracker.average_entry_price(), 0.0, "Entry price should be 0.0 after reset");
    
    info!("Test 2 PASSED: Portfolio resets correctly at episode boundaries");
}

#[test]
fn test_done_flag_correctness() {
    // Test 3: Verify done flag is true ONLY at episode boundaries
    //
    // done should be:
    // - false at samples 1, 2, ..., 199
    // - true at sample 200
    // - false at samples 201, 202, ..., 399
    // - true at sample 400
    // - etc.
    
    let dataset_size = 1_000; // Use smaller dataset for speed
    let expected_boundaries = 5; // 200, 400, 600, 800, 1000
    
    let mut done_count = 0;
    let mut last_done_index = 0;
    
    for i in 0..dataset_size {
        let done = (i + 1) % EPISODE_LENGTH == 0 || i + 1 >= dataset_size;
        
        if done {
            done_count += 1;
            
            // Verify spacing between boundaries (should be EPISODE_LENGTH)
            if done_count > 1 {
                let spacing = (i + 1) - last_done_index;
                assert_eq!(spacing, EPISODE_LENGTH, 
                    "Episode spacing should be {} samples, got {}", EPISODE_LENGTH, spacing);
            }
            
            last_done_index = i + 1;
        } else {
            // Verify done is false between boundaries
            let next_boundary = ((i / EPISODE_LENGTH) + 1) * EPISODE_LENGTH;
            assert!(i + 1 < next_boundary, 
                "Sample {} should have done=false (next boundary at {})", i + 1, next_boundary);
        }
    }
    
    assert_eq!(done_count, expected_boundaries, 
        "Expected {} terminal states, got {}", expected_boundaries, done_count);
        
    info!(done_count, "Test 3 PASSED: Done flags marked correctly");
}

#[test]
fn test_bellman_equation_respects_boundaries() {
    // Test 4: Verify Bellman equation doesn't bootstrap when done=true
    //
    // TD Target calculation:
    // - When done=false: TD_target = reward + γ × Q(next_state)
    // - When done=true:  TD_target = reward (no bootstrap)
    //
    // This test simulates the TD target calculation logic.
    
    let gamma = 0.99;
    let reward = 1.0;
    let next_q_value = 10.0;
    
    // Case 1: Mid-episode (done=false) - should bootstrap
    let done_false = false;
    let td_target_bootstrap = if done_false {
        reward
    } else {
        reward + gamma * next_q_value
    };
    
    assert_eq!(td_target_bootstrap, reward + gamma * next_q_value, 
        "Mid-episode should bootstrap from next Q-value");
    
    // Case 2: Episode boundary (done=true) - should NOT bootstrap
    let done_true = true;
    let td_target_terminal = if done_true {
        reward
    } else {
        reward + gamma * next_q_value
    };
    
    assert_eq!(td_target_terminal, reward, 
        "Episode boundary should use reward only (no bootstrap)");
    
    // Verify the difference
    let bootstrap_contribution: f64 = td_target_bootstrap - td_target_terminal;
    let expected_contribution: f64 = gamma * next_q_value;

    assert!((bootstrap_contribution - expected_contribution).abs() < 1e-6f64,
        "Bootstrap contribution should be γ × Q(next) = {}, got {}",
        expected_contribution, bootstrap_contribution);
        
    info!("Test 4 PASSED: Bellman equation respects episode boundaries");
}

#[test]
fn test_multi_epoch_episode_counting() {
    // Test 5: Verify episode counting across multiple epochs
    //
    // Each epoch should have:
    // - 90 episodes (18000 / 200 = 90)
    // - Episodes reset each epoch (not cumulative)
    
    let dataset_size = 18_000;
    let num_epochs = 3;
    let expected_episodes_per_epoch = 90;
    
    let mut total_episodes = 0;
    
    for epoch in 0..num_epochs {
        let mut epoch_episodes = 0;
        
        for i in 0..dataset_size {
            let done = (i + 1) % EPISODE_LENGTH == 0 || i + 1 >= dataset_size;
            if done {
                epoch_episodes += 1;
            }
        }
        
        assert_eq!(epoch_episodes, expected_episodes_per_epoch, 
            "Epoch {} should have {} episodes, got {}", 
            epoch + 1, expected_episodes_per_epoch, epoch_episodes);
        
        total_episodes += epoch_episodes;
    }
    
    assert_eq!(total_episodes, num_epochs * expected_episodes_per_epoch, 
        "Total episodes across {} epochs should be {}, got {}", 
        num_epochs, num_epochs * expected_episodes_per_epoch, total_episodes);
        
    info!(num_epochs, total_episodes, "Test 5 PASSED: Episode counting correct across epochs");
}

#[test]
fn test_partial_episode_at_epoch_end() {
    // Test 6: Edge case - Partial episode when dataset not divisible by EPISODE_LENGTH
    //
    // Example: 18,100 samples
    // - Episodes 1-90: 200 samples each (17,800 total)
    // - Episode 91: 300 samples (partial, 18,000-18,100)
    
    let dataset_size = 18_100;
    let expected_full_episodes = 90; // 18,000 / 200
    let expected_total_episodes = 91; // 90 full + 1 partial
    
    let mut boundary_count = 0;
    let mut episode_sizes = Vec::new();
    let mut last_boundary = 0;
    
    for i in 0..dataset_size {
        let done = (i + 1) % EPISODE_LENGTH == 0 || i + 1 >= dataset_size;
        if done {
            boundary_count += 1;
            let episode_size = (i + 1) - last_boundary;
            episode_sizes.push(episode_size);
            last_boundary = i + 1;
        }
    }
    
    assert_eq!(boundary_count, expected_total_episodes, 
        "Expected {} episodes (90 full + 1 partial), got {}", 
        expected_total_episodes, boundary_count);
    
    // Verify first 90 episodes are full size
    for i in 0..expected_full_episodes {
        assert_eq!(episode_sizes[i], EPISODE_LENGTH, 
            "Episode {} should be {} samples, got {}", 
            i + 1, EPISODE_LENGTH, episode_sizes[i]);
    }
    
    // Verify last episode is partial
    let last_episode_size = episode_sizes[episode_sizes.len() - 1];
    let expected_partial_size = dataset_size % EPISODE_LENGTH;
    assert_eq!(last_episode_size, expected_partial_size, 
        "Last episode should be {} samples (partial), got {}", 
        expected_partial_size, last_episode_size);
        
    info!(last_episode_size, "Test 6 PASSED: Partial episode handled correctly");
}

#[test]
fn test_episode_boundary_logging() {
    // Test 7: Verify episode boundary debug logging format
    //
    // Expected log format:
    // "Episode boundary at sample {i+1}/{total}, portfolio reset (episode {ep}/{total_ep})"
    
    let dataset_size = 1_000; // Small dataset for speed
    let _expected_boundaries = 5; // 200, 400, 600, 800, 1000

    for i in 0..dataset_size {
        let done = (i + 1) % EPISODE_LENGTH == 0 || i + 1 >= dataset_size;
        
        if done {
            let sample_num = i + 1;
            // Episode number: For sample 200 → episode 1, sample 400 → episode 2, etc.
            let episode_num = if sample_num % EPISODE_LENGTH == 0 {
                sample_num / EPISODE_LENGTH
            } else {
                (sample_num / EPISODE_LENGTH) + 1
            };
            let _total_episodes = (dataset_size + EPISODE_LENGTH - 1) / EPISODE_LENGTH;

            // Verify episode numbering matches the logic used in dqn.rs
            let expected_episode = if sample_num % EPISODE_LENGTH == 0 {
                sample_num / EPISODE_LENGTH
            } else {
                (sample_num / EPISODE_LENGTH) + 1
            };

            assert_eq!(episode_num, expected_episode,
                "Episode number mismatch at sample {}: expected {}, got {}",
                sample_num, expected_episode, episode_num);
        }
    }
    
    info!("Test 7 PASSED: Episode boundary logging format correct");
}

#[test]
fn test_short_dataset_edge_case() {
    // Test 8: Edge case - Dataset smaller than EPISODE_LENGTH
    //
    // Example: 50 samples with EPISODE_LENGTH=200
    // Expected: 1 episode (partial) with terminal state at sample 50
    
    let dataset_size = 50;
    let expected_episodes = 1;
    
    let mut boundary_count = 0;
    
    for i in 0..dataset_size {
        let done = (i + 1) % EPISODE_LENGTH == 0 || i + 1 >= dataset_size;
        if done {
            boundary_count += 1;
            
            // Verify this is the last sample
            assert_eq!(i + 1, dataset_size, 
                "Short dataset should have terminal state only at last sample");
        }
    }
    
    assert_eq!(boundary_count, expected_episodes, 
        "Short dataset ({} samples) should have {} episode, got {}", 
        dataset_size, expected_episodes, boundary_count);
        
    info!("Test 8 PASSED: Short dataset edge case handled correctly");
}

#[test]
fn test_exact_multiple_dataset() {
    // Test 9: Edge case - Dataset size exactly divisible by EPISODE_LENGTH
    //
    // Example: 1,000 samples with EPISODE_LENGTH=200
    // Expected: 5 episodes (200, 400, 600, 800, 1000)
    
    let dataset_size = 1_000; // Exactly 5 episodes
    let expected_episodes = 5;
    
    let mut boundary_count = 0;
    let mut boundary_indices = Vec::new();
    
    for i in 0..dataset_size {
        let done = (i + 1) % EPISODE_LENGTH == 0 || i + 1 >= dataset_size;
        if done {
            boundary_count += 1;
            boundary_indices.push(i + 1);
        }
    }
    
    assert_eq!(boundary_count, expected_episodes, 
        "Dataset of {} samples should have exactly {} episodes", 
        dataset_size, expected_episodes);
    
    // Verify all boundaries are at exact multiples
    for (idx, &boundary) in boundary_indices.iter().enumerate() {
        let expected_boundary = (idx + 1) * EPISODE_LENGTH;
        assert_eq!(boundary, expected_boundary, 
            "Episode {} boundary should be at sample {}, got {}", 
            idx + 1, expected_boundary, boundary);
    }
    
    info!("Test 9 PASSED: Exact multiple dataset handled correctly");
}

#[test]
fn test_no_position_leakage_between_episodes() {
    // Test 10: Verify no position state leaks between episodes
    //
    // Simulates 2 episodes:
    // - Episode 1: Open long position, end with done=true
    // - Episode 2: Should start with clean portfolio (no position)
    
    let mut tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);

    // Episode 1: Open position
    tracker.execute_legacy_action(TradingAction::Buy, 100.0, 10.0);
    assert_eq!(tracker.current_position(), 10.0, "Episode 1 should have position");
    
    // Episode boundary: Reset portfolio
    tracker.reset();
    
    // Episode 2: Verify clean start
    assert_eq!(tracker.current_position(), 0.0, "Episode 2 should start flat");
    assert_eq!(tracker.cash_balance(), 10_000.0, "Episode 2 should have initial capital");
    assert_eq!(tracker.realized_pnl(), 0.0, "Episode 2 should have zero realized PnL");
    
    // Verify unrealized PnL is also zero
    let unrealized = tracker.unrealized_pnl(100.0);
    assert_eq!(unrealized, 0.0, "Episode 2 should have zero unrealized PnL");
    
    info!("Test 10 PASSED: No position leakage between episodes");
}

// Integration test helper to count terminal states in a mock training loop
fn count_terminal_states_in_epoch(dataset_size: usize) -> usize {
    let mut terminal_count = 0;
    
    for i in 0..dataset_size {
        let done = (i + 1) % EPISODE_LENGTH == 0 || i + 1 >= dataset_size;
        if done {
            terminal_count += 1;
        }
    }
    
    terminal_count
}

#[test]
fn test_terminal_state_count_18k_dataset() {
    // Test 11: Comprehensive test - 18,000 sample dataset
    //
    // Verify:
    // - 90 terminal states per epoch (18000 / 200 = 90)
    // - Terminal states at samples: 200, 400, ..., 17800, 18000
    
    let dataset_size = 18_000;
    let expected_terminals = 90;
    
    let terminal_count = count_terminal_states_in_epoch(dataset_size);
    
    assert_eq!(terminal_count, expected_terminals, 
        "18K dataset should have {} terminal states, got {}", 
        expected_terminals, terminal_count);
    
    // Calculate improvement over old implementation
    let old_terminals = 1; // Before fix: only final sample
    let improvement_factor = terminal_count / old_terminals;
    
    info!(terminal_count, old_terminals, improvement_factor, "Test 11 PASSED: Terminal state count correct for 18K dataset");
}

#[test]
fn test_episode_boundary_consistency() {
    // Test 12: Verify episode boundaries are consistent across iterations
    //
    // Run boundary detection 100 times, verify same results
    
    let dataset_size = 1_000;
    let iterations = 100;
    
    let mut all_boundary_counts = Vec::new();
    
    for _ in 0..iterations {
        let count = count_terminal_states_in_epoch(dataset_size);
        all_boundary_counts.push(count);
    }
    
    // Verify all counts are identical
    let first_count = all_boundary_counts[0];
    for &count in &all_boundary_counts {
        assert_eq!(count, first_count, 
            "Episode boundary count should be consistent across iterations");
    }
    
    info!(iterations, "Test 12 PASSED: Episode boundaries consistent across iterations");
}