#![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,
)]
//! TDD Tests for Learning Rate Scheduler with Warmup
//!
//! Tests cover:
//! - Linear warmup period
//! - Constant learning rate
//! - Linear decay
//! - Cosine annealing decay
//! - Edge cases (warmup_steps=0, step boundaries)

use ml::trainers::dqn::lr_scheduler::{LRDecayType, LRScheduler};

#[test]
fn test_lr_scheduler_constant() {
    let scheduler = LRScheduler::new(0.001, 0, LRDecayType::Constant);

    assert_eq!(scheduler.get_lr(), 0.001);

    // Step forward and verify LR stays constant
    let mut scheduler = scheduler;
    scheduler.step();
    assert_eq!(scheduler.get_lr(), 0.001);

    scheduler.step();
    assert_eq!(scheduler.get_lr(), 0.001);

    for _ in 0..100 {
        scheduler.step();
    }
    assert_eq!(scheduler.get_lr(), 0.001);
}

#[test]
fn test_lr_scheduler_linear_warmup() {
    let initial_lr = 0.001;
    let warmup_steps = 10;
    let scheduler = LRScheduler::new(initial_lr, warmup_steps, LRDecayType::Constant);

    // Step 0: 0% warmup
    assert_eq!(scheduler.get_lr(), 0.0);

    // Step 1: 10% warmup
    let mut scheduler = scheduler;
    scheduler.step();
    assert!((scheduler.get_lr() - 0.0001).abs() < 1e-6);

    // Step 5: 50% warmup
    for _ in 0..4 {
        scheduler.step();
    }
    assert!((scheduler.get_lr() - 0.0005).abs() < 1e-6);

    // Step 10: 100% warmup (full LR)
    for _ in 0..5 {
        scheduler.step();
    }
    assert!((scheduler.get_lr() - 0.001).abs() < 1e-6);

    // Step 11+: stays at full LR (constant decay mode)
    scheduler.step();
    assert_eq!(scheduler.get_lr(), 0.001);
}

#[test]
fn test_lr_scheduler_linear_decay() {
    let initial_lr = 0.001;
    let warmup_steps = 5;
    let total_steps = 100;
    let end_lr = 0.0001;

    let scheduler = LRScheduler::new(
        initial_lr,
        warmup_steps,
        LRDecayType::Linear { end_lr, total_steps },
    );

    // Step 0: warmup (0%)
    assert_eq!(scheduler.get_lr(), 0.0);

    // Step 5: end of warmup (100%)
    let mut scheduler = scheduler;
    for _ in 0..5 {
        scheduler.step();
    }
    assert!((scheduler.get_lr() - 0.001).abs() < 1e-6);

    // Step 50: halfway through decay
    for _ in 0..45 {
        scheduler.step();
    }
    let expected_lr = 0.001 - (0.001 - 0.0001) * (50 - 5) as f64 / (100 - 5) as f64;
    assert!((scheduler.get_lr() - expected_lr).abs() < 1e-6);

    // Step 100: end of decay
    for _ in 0..50 {
        scheduler.step();
    }
    assert!((scheduler.get_lr() - 0.0001).abs() < 1e-6);

    // Step 101+: stays at end_lr
    scheduler.step();
    assert!((scheduler.get_lr() - 0.0001).abs() < 1e-6);
}

#[test]
fn test_lr_scheduler_cosine_decay() {
    let initial_lr = 0.001;
    let warmup_steps = 5;
    let total_steps = 100;
    let min_lr = 0.0001;

    let scheduler = LRScheduler::new(
        initial_lr,
        warmup_steps,
        LRDecayType::Cosine { min_lr, total_steps },
    );

    // Step 0: warmup (0%)
    assert_eq!(scheduler.get_lr(), 0.0);

    // Step 5: end of warmup (100%)
    let mut scheduler = scheduler;
    for _ in 0..5 {
        scheduler.step();
    }
    assert!((scheduler.get_lr() - 0.001).abs() < 1e-6);

    // Step 6: start of cosine decay
    scheduler.step();
    let lr_after_warmup = scheduler.get_lr();
    assert!(lr_after_warmup < 0.001); // Should start decaying
    assert!(lr_after_warmup > 0.0001); // But still above min

    // Step 100: end of decay (should be at min_lr)
    for _ in 0..94 {
        scheduler.step();
    }
    assert!((scheduler.get_lr() - 0.0001).abs() < 1e-5);

    // Step 101+: stays at min_lr
    scheduler.step();
    assert!((scheduler.get_lr() - 0.0001).abs() < 1e-5);
}

#[test]
fn test_lr_scheduler_no_warmup() {
    let initial_lr = 0.001;
    let warmup_steps = 0;
    let total_steps = 50;
    let end_lr = 0.0001;

    let scheduler = LRScheduler::new(
        initial_lr,
        warmup_steps,
        LRDecayType::Linear { end_lr, total_steps },
    );

    // Step 0: no warmup, immediately at initial_lr
    assert_eq!(scheduler.get_lr(), 0.001);

    // Step 25: halfway through decay
    let mut scheduler = scheduler;
    for _ in 0..25 {
        scheduler.step();
    }
    let expected_lr = 0.001 - (0.001 - 0.0001) * 0.5;
    assert!((scheduler.get_lr() - expected_lr).abs() < 1e-6);
}

#[test]
fn test_lr_scheduler_cosine_shape() {
    let initial_lr = 1.0;
    let warmup_steps = 0;
    let total_steps = 100;
    let min_lr = 0.0;

    let mut scheduler = LRScheduler::new(
        initial_lr,
        warmup_steps,
        LRDecayType::Cosine { min_lr, total_steps },
    );

    // Step 0: initial_lr
    assert_eq!(scheduler.get_lr(), 1.0);

    // Step 25: ~75% of initial_lr (cosine is slow at start)
    for _ in 0..25 {
        scheduler.step();
    }
    let lr_25 = scheduler.get_lr();
    assert!(lr_25 > 0.7 && lr_25 < 0.9);

    // Step 50: ~50% of initial_lr (cosine at midpoint)
    for _ in 0..25 {
        scheduler.step();
    }
    let lr_50 = scheduler.get_lr();
    assert!(lr_50 > 0.4 && lr_50 < 0.6);

    // Step 75: ~25% of initial_lr (cosine is fast at end)
    for _ in 0..25 {
        scheduler.step();
    }
    let lr_75 = scheduler.get_lr();
    assert!(lr_75 > 0.1 && lr_75 < 0.3);

    // Step 100: min_lr
    for _ in 0..25 {
        scheduler.step();
    }
    assert!(scheduler.get_lr().abs() < 1e-5);
}

#[test]
fn test_lr_scheduler_reset() {
    let initial_lr = 0.001;
    let warmup_steps = 5;
    let mut scheduler = LRScheduler::new(initial_lr, warmup_steps, LRDecayType::Constant);

    // Step forward
    for _ in 0..10 {
        scheduler.step();
    }
    assert_eq!(scheduler.get_lr(), 0.001);

    // Reset
    scheduler.reset();
    assert_eq!(scheduler.get_current_step(), 0);
    assert_eq!(scheduler.get_lr(), 0.0); // Back to warmup start
}

#[test]
fn test_lr_scheduler_get_current_step() {
    let mut scheduler = LRScheduler::new(0.001, 0, LRDecayType::Constant);

    assert_eq!(scheduler.get_current_step(), 0);

    scheduler.step();
    assert_eq!(scheduler.get_current_step(), 1);

    for _ in 0..99 {
        scheduler.step();
    }
    assert_eq!(scheduler.get_current_step(), 100);
}