//! Tests for DQN network initialization randomness
//!
//! Validates that Xavier initialization produces different random weights
//! for each network instance (not deterministic/identical weights).

// No unused imports
use ml::dqn::dqn::{WorkingDQN, WorkingDQNConfig};
use ml::MLError;

/// Extract weight values from the first layer of a DQN network
fn extract_first_layer_weights(agent: &WorkingDQN) -> Result<Vec<f32>, MLError> {
    // Get weights from VarMap
    let vars = agent.get_q_network_vars();
    let vars_data = vars.data().lock().unwrap();

    // Find first hidden layer weights
    let (name, var) = vars_data
        .iter()
        .find(|(name, _)| name.contains("hidden_0") && name.contains("weight"))
        .ok_or_else(|| {
            MLError::ModelError("Failed to find hidden_0.weight in VarMap".to_string())
        })?;

    println!("Found weight tensor: {}", name);

    // Get tensor and convert to Vec<f32>
    let weight_tensor = var.as_tensor();
    let weight_vec = weight_tensor
        .flatten_all()
        .map_err(|e| MLError::ModelError(format!("Failed to flatten weights: {}", e)))?
        .to_vec1::<f32>()
        .map_err(|e| MLError::ModelError(format!("Failed to convert weights to vec: {}", e)))?;

    Ok(weight_vec)
}

/// Test that multiple DQN instances have different initial weights
#[test]
fn test_initialization_randomness() -> Result<(), MLError> {
    // Create 5 DQN instances with identical configuration
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 128;
    config.hidden_dims = vec![64, 32];
    config.num_actions = 3;

    let agents: Vec<WorkingDQN> = (0..5)
        .map(|_| WorkingDQN::new(config.clone()))
        .collect::<Result<Vec<_>, _>>()?;

    println!(
        "Created {} DQN instances with identical config",
        agents.len()
    );

    // Extract initial weights from each agent
    let all_weights: Vec<Vec<f32>> = agents
        .iter()
        .enumerate()
        .map(|(i, agent)| {
            extract_first_layer_weights(agent).map(|w| {
                println!("Agent {}: Extracted {} weights", i, w.len());
                w
            })
        })
        .collect::<Result<Vec<_>, _>>()?;

    // Verify all weight vectors have the same length
    let weight_len = all_weights[0].len();
    for (i, weights) in all_weights.iter().enumerate() {
        assert_eq!(
            weights.len(),
            weight_len,
            "Agent {} has different weight vector length",
            i
        );
    }

    println!("All agents have {} weights in first layer", weight_len);

    // Compare each pair of agents
    let mut num_identical_pairs = 0;
    let mut max_similarity = 0.0f32;

    for i in 0..5 {
        for j in (i + 1)..5 {
            let weights_i = &all_weights[i];
            let weights_j = &all_weights[j];

            // Calculate element-wise absolute differences
            let diffs: Vec<f32> = weights_i
                .iter()
                .zip(weights_j.iter())
                .map(|(a, b)| (a - b).abs())
                .collect();

            let mean_diff = diffs.iter().sum::<f32>() / diffs.len() as f32;
            let max_diff = diffs.iter().cloned().fold(0.0f32, f32::max);

            // Calculate similarity (1.0 = identical, 0.0 = completely different)
            let similarity = 1.0 - mean_diff.min(1.0);
            max_similarity = max_similarity.max(similarity);

            println!(
                "Agent {} vs Agent {}: mean_diff={:.6}, max_diff={:.6}, similarity={:.4}",
                i, j, mean_diff, max_diff, similarity
            );

            // Check if weights are identical (within floating point tolerance)
            let are_identical = mean_diff < 1e-6;
            if are_identical {
                num_identical_pairs += 1;
            }

            // Verify weights are DIFFERENT (not identical)
            assert!(
                !are_identical,
                "Agent {} and Agent {} have IDENTICAL weights (mean_diff={:.8})",
                i, j, mean_diff
            );
        }
    }

    println!("\nRandomness verification:");
    println!("  Identical pairs: {} / 10", num_identical_pairs);
    println!("  Max similarity: {:.4}", max_similarity);

    assert_eq!(
        num_identical_pairs, 0,
        "Expected 0 identical pairs, found {}",
        num_identical_pairs
    );

    Ok(())
}

/// Test that initialization variance is within expected Xavier bounds
#[test]
fn test_initialization_variance() -> Result<(), MLError> {
    // Create 5 DQN instances
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 128;
    config.hidden_dims = vec![64, 32];
    config.num_actions = 3;

    let agents: Vec<WorkingDQN> = (0..5)
        .map(|_| WorkingDQN::new(config.clone()))
        .collect::<Result<Vec<_>, _>>()?;

    // Extract weights from all agents
    let all_weights: Vec<Vec<f32>> = agents
        .iter()
        .map(|agent| extract_first_layer_weights(agent))
        .collect::<Result<Vec<_>, _>>()?;

    // Calculate variance across all weights from all agents
    let weight_len = all_weights[0].len();
    let mut variances = Vec::with_capacity(weight_len);

    for idx in 0..weight_len {
        // Collect weight at position `idx` from all agents
        let values: Vec<f32> = all_weights.iter().map(|w| w[idx]).collect();

        // Calculate variance: var = mean((x - mean(x))^2)
        let mean = values.iter().sum::<f32>() / values.len() as f32;
        let variance = values
            .iter()
            .map(|v| {
                let diff = v - mean;
                diff * diff
            })
            .sum::<f32>()
            / values.len() as f32;

        variances.push(variance);
    }

    // Calculate mean variance across all weight positions
    let mean_variance = variances.iter().sum::<f32>() / variances.len() as f32;
    let max_variance = variances.iter().cloned().fold(0.0f32, f32::max);
    let min_variance = variances.iter().cloned().fold(f32::MAX, f32::min);

    println!("Initialization variance statistics:");
    println!("  Mean variance: {:.6}", mean_variance);
    println!("  Max variance:  {:.6}", max_variance);
    println!("  Min variance:  {:.6}", min_variance);

    // Xavier initialization for input_dim=128, output_dim=64:
    // Theoretical variance = 2 / (128 + 64) = 0.0104
    let expected_variance = 2.0 / (128.0 + 64.0);
    println!("  Expected (Xavier): {:.6}", expected_variance);

    // Verify variance is non-zero (randomness exists)
    assert!(
        mean_variance > 0.001,
        "Variance too low ({:.6}), weights may not be random",
        mean_variance
    );

    // Verify variance is within reasonable bounds for Xavier initialization
    // Allow 10x range due to sampling variance with only 5 agents
    assert!(
        mean_variance < expected_variance * 10.0,
        "Variance too high ({:.6}), exceeds Xavier expected range ({:.6})",
        mean_variance,
        expected_variance * 10.0
    );

    Ok(())
}

/// Test that different layers have different random weights
#[test]
fn test_layer_independence() -> Result<(), MLError> {
    // Create one agent
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 128;
    config.hidden_dims = vec![64, 32];
    config.num_actions = 3;

    let agent = WorkingDQN::new(config)?;

    // Extract weights from multiple layers
    let vars = agent.get_q_network_vars();
    let vars_data = vars.data().lock().unwrap();

    let hidden0_weights = vars_data
        .iter()
        .find(|(name, _)| name.contains("hidden_0") && name.contains("weight"))
        .ok_or_else(|| MLError::ModelError("Missing hidden_0.weight".to_string()))?
        .1
        .as_tensor()
        .flatten_all()?
        .to_vec1::<f32>()?;

    let hidden1_weights = vars_data
        .iter()
        .find(|(name, _)| name.contains("hidden_1") && name.contains("weight"))
        .ok_or_else(|| MLError::ModelError("Missing hidden_1.weight".to_string()))?
        .1
        .as_tensor()
        .flatten_all()?
        .to_vec1::<f32>()?;

    println!("Hidden layer 0: {} weights", hidden0_weights.len());
    println!("Hidden layer 1: {} weights", hidden1_weights.len());

    // Compare first N weights (where N = min layer size)
    let compare_len = hidden0_weights.len().min(hidden1_weights.len());

    let mut identical_count = 0;
    for i in 0..compare_len {
        if (hidden0_weights[i] - hidden1_weights[i]).abs() < 1e-6 {
            identical_count += 1;
        }
    }

    let identical_ratio = identical_count as f32 / compare_len as f32;
    println!(
        "Identical weights: {} / {} ({:.2}%)",
        identical_count,
        compare_len,
        identical_ratio * 100.0
    );

    // Verify layers have different weights (not copied)
    // Allow up to 5% coincidental matches due to randomness
    assert!(
        identical_ratio < 0.05,
        "Layers have {:.2}% identical weights (expected <5%)",
        identical_ratio * 100.0
    );

    Ok(())
}

/// Test that target network is properly copied from main network
#[test]
fn test_target_network_copy() -> Result<(), MLError> {
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.state_dim = 128;
    config.hidden_dims = vec![64, 32];
    config.num_actions = 3;

    let agent = WorkingDQN::new(config)?;

    // Extract weights from both networks
    let main_vars = agent.get_q_network_vars();
    let main_data = main_vars.data().lock().unwrap();
    let main_weights = main_data
        .iter()
        .find(|(name, _)| name.contains("hidden_0") && name.contains("weight"))
        .ok_or_else(|| MLError::ModelError("Missing main hidden_0.weight".to_string()))?
        .1
        .as_tensor()
        .flatten_all()?
        .to_vec1::<f32>()?;
    drop(main_data);

    let target_vars = agent.get_target_network_vars();
    let target_data = target_vars.data().lock().unwrap();
    let target_weights = target_data
        .iter()
        .find(|(name, _)| name.contains("hidden_0") && name.contains("weight"))
        .ok_or_else(|| MLError::ModelError("Missing target hidden_0.weight".to_string()))?
        .1
        .as_tensor()
        .flatten_all()?
        .to_vec1::<f32>()?;
    drop(target_data);

    println!("Main network:   {} weights", main_weights.len());
    println!("Target network: {} weights", target_weights.len());

    // Verify both networks have the same number of weights
    assert_eq!(
        main_weights.len(),
        target_weights.len(),
        "Main and target networks have different sizes"
    );

    // Calculate mean absolute difference
    let mean_diff: f32 = main_weights
        .iter()
        .zip(target_weights.iter())
        .map(|(a, b)| (a - b).abs())
        .sum::<f32>()
        / main_weights.len() as f32;

    println!("Mean absolute difference: {:.8}", mean_diff);

    // Verify target network is initialized as a copy of main network
    // (weights should be identical within floating point tolerance)
    assert!(
        mean_diff < 1e-6,
        "Target network weights differ from main network (mean_diff={:.8})",
        mean_diff
    );

    Ok(())
}