foxhunt/ml/tests/target_network_weight_verification_test.rs

//! Target Network Weight Verification Tests
//!
//! CRITICAL: The existing target network tests only verify UPDATE COUNTERS,
//! they NEVER verify that weights are actually copied. This is CATASTROPHIC
//! untested code that could be silently broken.
//!
//! These tests verify:
//! 1. Hard updates actually copy weights (element-wise equality)
//! 2. Soft updates (Polyak averaging) correctly blend weights
//! 3. Weights diverge during training before update
//! 4. Weight magnitude checks (no NaN/Inf corruption)

use anyhow::Result;
use ml::dqn::dqn::{WorkingDQN, WorkingDQNConfig};
use ml::dqn::Experience;

/// Helper: Create test DQN with configurable soft/hard update mode
fn create_test_dqn(use_soft_updates: bool, tau: f64, target_update_freq: usize) -> Result<WorkingDQN> {
    let config = WorkingDQNConfig {
        state_dim: 10,
        num_actions: 3,
        hidden_dims: vec![16],
        learning_rate: 0.001,
        gamma: 0.99,
        epsilon_start: 0.0, // Disable exploration for deterministic tests
        epsilon_end: 0.0,
        epsilon_decay: 1.0,
        replay_buffer_capacity: 1000,
        batch_size: 4,
        min_replay_size: 4,
        target_update_freq,
        use_double_dqn: false,
        use_huber_loss: true,
        huber_delta: 1.0,
        leaky_relu_alpha: 0.01,
        gradient_clip_norm: 10.0,
        tau,
        use_soft_updates,
        warmup_steps: 0, // No warmup for tests
        temperature_start: 1.0,
        temperature_decay: 0.99,
    };

    WorkingDQN::new(config).map_err(|e| anyhow::anyhow!("DQN creation failed: {}", e))
}

/// Helper: Add experiences to replay buffer
fn populate_replay_buffer(dqn: &WorkingDQN, count: usize) -> Result<()> {
    for i in 0..count {
        let state = vec![i as f32 * 0.1; 10];
        let next_state = vec![(i + 1) as f32 * 0.1; 10];
        let action = (i % 3) as u8;
        let reward = i as f32;
        let done = false;

        let experience = Experience::new(state, action, reward, next_state, done);
        dqn.store_experience(experience)
            .map_err(|e| anyhow::anyhow!("Failed to store experience: {}", e))?;
    }
    Ok(())
}

/// Helper: Extract all weights from a VarMap as a flat vector (sorted by name for consistency)
fn extract_weights(vars: &candle_nn::VarMap) -> Result<Vec<f32>> {
    let vars_data = vars
        .data()
        .lock()
        .map_err(|e| anyhow::anyhow!("Failed to lock vars: {}", e))?;

    // Sort by name to ensure consistent ordering
    let mut sorted_vars: Vec<_> = vars_data.iter().collect();
    sorted_vars.sort_by_key(|(name, _)| *name);

    let mut weights = Vec::new();
    for (_name, var) in sorted_vars {
        let tensor = var.as_tensor();

        // Handle different tensor shapes (1D weights, 2D bias)
        let data = tensor
            .flatten_all()
            .map_err(|e| anyhow::anyhow!("Failed to flatten tensor: {}", e))?
            .to_vec1::<f32>()
            .map_err(|e| anyhow::anyhow!("Failed to extract tensor data: {}", e))?;

        weights.extend(data);
    }
    Ok(weights)
}

/// Helper: Compare two weight vectors with tolerance
fn weights_equal(a: &[f32], b: &[f32], tolerance: f32) -> bool {
    if a.len() != b.len() {
        return false;
    }

    a.iter()
        .zip(b.iter())
        .all(|(x, y)| (x - y).abs() < tolerance)
}

/// Test 1: CRITICAL - Hard update actually copies weights (element-wise equality)
#[test]
fn test_hard_update_actually_copies_weights() -> Result<()> {
    let mut dqn = create_test_dqn(false, 1.0, 10)?; // Hard updates every 10 steps
    populate_replay_buffer(&dqn, 20)?;

    // Train for 9 steps (no update yet)
    for _ in 0..9 {
        let _ = dqn.train_step(None);
    }

    // Get online network weights before update
    let online_weights_before_update = extract_weights(dqn.get_q_network_vars())?;
    let target_weights_before_update = extract_weights(dqn.get_target_network_vars())?;

    // Verify weights have diverged during training
    assert!(
        !weights_equal(&online_weights_before_update, &target_weights_before_update, 1e-6),
        "Online and target networks should diverge during training"
    );

    // Step 10: Trigger hard update
    let _ = dqn.train_step(None);

    // Get weights after hard update
    let online_weights_after = extract_weights(dqn.get_q_network_vars())?;
    let target_weights_after = extract_weights(dqn.get_target_network_vars())?;

    // CRITICAL ASSERTION: After hard update, target == online (element-wise)
    assert!(
        weights_equal(&target_weights_after, &online_weights_after, 1e-6),
        "CATASTROPHIC BUG: Hard update did NOT copy weights!\n\
         Target != Online after hard update at step 10.\n\
         Expected: target_weights == online_weights\n\
         Actual: {} weight mismatches detected",
        target_weights_after
            .iter()
            .zip(online_weights_after.iter())
            .filter(|(&a, &b)| (a - b).abs() >= 1e-6)
            .count()
    );

    println!("✓ Hard update verified: {} weights copied correctly", target_weights_after.len());
    Ok(())
}

/// Test 2: CRITICAL - Soft update (Polyak averaging) blends weights correctly
#[test]
fn test_polyak_update_blends_weights() -> Result<()> {
    let tau = 0.5; // 50% blending for easy verification
    let mut dqn = create_test_dqn(true, tau, 10)?; // Soft updates (but freq doesn't matter)
    populate_replay_buffer(&dqn, 20)?;

    // Get initial weights (should be equal at initialization)
    let initial_online = extract_weights(dqn.get_q_network_vars())?;
    let initial_target = extract_weights(dqn.get_target_network_vars())?;

    // Verify initialization: target == online
    assert!(
        weights_equal(&initial_online, &initial_target, 1e-6),
        "Initial weights should be equal"
    );

    // Train for 1 step (soft update happens every step)
    let _ = dqn.train_step(None);

    // Get updated weights
    let online_after = extract_weights(dqn.get_q_network_vars())?;
    let target_after = extract_weights(dqn.get_target_network_vars())?;

    // Verify online network changed during training
    assert!(
        !weights_equal(&initial_online, &online_after, 1e-6),
        "Online network should change during training"
    );

    // CRITICAL VERIFICATION: target_new = tau * online + (1-tau) * target_old
    // With tau=0.5: target_new = 0.5 * online_after + 0.5 * initial_target
    let mut polyak_correct_count = 0;
    let mut polyak_wrong_count = 0;

    for i in 0..target_after.len() {
        let expected = tau as f32 * online_after[i] + (1.0 - tau as f32) * initial_target[i];
        let actual = target_after[i];
        let error = (expected - actual).abs();

        if error < 1e-4 {
            polyak_correct_count += 1;
        } else {
            polyak_wrong_count += 1;
            if polyak_wrong_count <= 5 {
                println!(
                    "Weight[{}]: expected={:.6}, actual={:.6}, error={:.6}",
                    i, expected, actual, error
                );
            }
        }
    }

    assert!(
        polyak_correct_count > polyak_wrong_count,
        "CATASTROPHIC BUG: Polyak averaging NOT implemented correctly!\n\
         Expected: target_new = {:.2} * online + {:.2} * target_old\n\
         Correct weights: {}/{}\n\
         Wrong weights: {}/{}",
        tau,
        1.0 - tau,
        polyak_correct_count,
        target_after.len(),
        polyak_wrong_count,
        target_after.len()
    );

    println!(
        "✓ Polyak averaging verified: {}/{} weights correct (tau={:.2})",
        polyak_correct_count,
        target_after.len(),
        tau
    );
    Ok(())
}

/// Test 3: Weights diverge during training (before update)
#[test]
fn test_weights_diverge_during_training() -> Result<()> {
    let mut dqn = create_test_dqn(false, 1.0, 100)?; // Hard update every 100 steps
    populate_replay_buffer(&dqn, 20)?;

    let initial_online = extract_weights(dqn.get_q_network_vars())?;
    let initial_target = extract_weights(dqn.get_target_network_vars())?;

    // Verify initial equality
    assert!(
        weights_equal(&initial_online, &initial_target, 1e-6),
        "Initial weights should be equal"
    );

    // Train for 50 steps (no hard update yet)
    for _ in 0..50 {
        let _ = dqn.train_step(None);
    }

    let online_after_50 = extract_weights(dqn.get_q_network_vars())?;
    let target_after_50 = extract_weights(dqn.get_target_network_vars())?;

    // Online should change
    assert!(
        !weights_equal(&initial_online, &online_after_50, 1e-6),
        "Online network should change during training"
    );

    // Target should NOT change (no update yet)
    assert!(
        weights_equal(&initial_target, &target_after_50, 1e-6),
        "Target network should NOT change before hard update"
    );

    // Online and target should be different
    let divergence_count = online_after_50
        .iter()
        .zip(target_after_50.iter())
        .filter(|(&a, &b)| (a - b).abs() >= 1e-6)
        .count();

    assert!(
        divergence_count > 0,
        "Networks should diverge during training (found {} divergent weights)",
        divergence_count
    );

    println!(
        "✓ Weight divergence verified: {}/{} weights diverged during training",
        divergence_count,
        online_after_50.len()
    );
    Ok(())
}

/// Test 4: No NaN/Inf corruption in target network weights
#[test]
fn test_no_nan_inf_in_target_weights() -> Result<()> {
    let mut dqn = create_test_dqn(false, 1.0, 10)?;
    populate_replay_buffer(&dqn, 20)?;

    // Train for 100 steps (multiple hard updates)
    for _ in 0..100 {
        let _ = dqn.train_step(None);
    }

    let target_weights = extract_weights(dqn.get_target_network_vars())?;

    let nan_count = target_weights.iter().filter(|x| x.is_nan()).count();
    let inf_count = target_weights.iter().filter(|x| x.is_infinite()).count();

    assert_eq!(
        nan_count, 0,
        "CATASTROPHIC: {} NaN values in target network weights!",
        nan_count
    );

    assert_eq!(
        inf_count, 0,
        "CATASTROPHIC: {} Inf values in target network weights!",
        inf_count
    );

    println!(
        "✓ Weight integrity verified: {}/{} weights valid (no NaN/Inf)",
        target_weights.len() - nan_count - inf_count,
        target_weights.len()
    );
    Ok(())
}

/// Test 5: Hard update resets divergence to zero
#[test]
fn test_hard_update_resets_divergence() -> Result<()> {
    let mut dqn = create_test_dqn(false, 1.0, 10)?;
    populate_replay_buffer(&dqn, 20)?;

    // Train for 9 steps (diverge)
    for _ in 0..9 {
        let _ = dqn.train_step(None);
    }

    let online_before = extract_weights(dqn.get_q_network_vars())?;
    let target_before = extract_weights(dqn.get_target_network_vars())?;

    let divergence_before = online_before
        .iter()
        .zip(target_before.iter())
        .map(|(a, b)| (a - b).abs())
        .sum::<f32>();

    assert!(
        divergence_before > 1e-4,
        "Networks should diverge before update"
    );

    // Step 10: Hard update
    let _ = dqn.train_step(None);

    let online_after = extract_weights(dqn.get_q_network_vars())?;
    let target_after = extract_weights(dqn.get_target_network_vars())?;

    let divergence_after = online_after
        .iter()
        .zip(target_after.iter())
        .map(|(a, b)| (a - b).abs())
        .sum::<f32>();

    assert!(
        divergence_after < 1e-4,
        "CATASTROPHIC: Hard update did NOT reset divergence!\n\
         Divergence before update: {:.6}\n\
         Divergence after update: {:.6}",
        divergence_before,
        divergence_after
    );

    println!(
        "✓ Divergence reset verified: {:.6} → {:.6}",
        divergence_before, divergence_after
    );
    Ok(())
}

/// Test 6: Soft update preserves target network stability (gradual change)
#[test]
fn test_soft_update_gradual_change() -> Result<()> {
    let tau = 0.001; // Very small tau for stability
    let mut dqn = create_test_dqn(true, tau, 10)?;
    populate_replay_buffer(&dqn, 20)?;

    let initial_target = extract_weights(dqn.get_target_network_vars())?;

    // Train for 10 steps (10 soft updates)
    for _ in 0..10 {
        let _ = dqn.train_step(None);
    }

    let target_after_10 = extract_weights(dqn.get_target_network_vars())?;

    // Calculate total change magnitude
    let total_change = initial_target
        .iter()
        .zip(target_after_10.iter())
        .map(|(a, b)| (a - b).abs())
        .sum::<f32>();

    let avg_change = total_change / initial_target.len() as f32;

    // With tau=0.001, target should change very slowly
    assert!(
        avg_change < 0.1,
        "Soft update should preserve stability (avg_change={:.6})",
        avg_change
    );

    // But target should NOT be identical (some change should occur)
    assert!(
        avg_change > 1e-6,
        "Soft update should cause SOME change (avg_change={:.6})",
        avg_change
    );

    println!(
        "✓ Soft update stability verified: avg_change={:.6} (tau={:.3})",
        avg_change, tau
    );
    Ok(())
}