foxhunt/ml/examples/diagnose_factored_network.rs

//! Diagnostic Tool: FactoredQNetwork Q-Value Uniqueness Analysis
//!
//! Verifies that FactoredQNetwork outputs 45 truly unique Q-values,
//! not repeating the same 8 values.
//!
//! Tests:
//! 1. Q-value uniqueness (count unique values per forward pass)
//! 2. Additive factorization formula correctness
//! 3. Value duplication analysis across action space
//! 4. Distribution of Q-values

use candle_core::{Device, Tensor};
use ml::dqn::factored_q_network::FactoredQNetwork;
use ml::MLError;
use std::collections::HashSet;

fn main() -> Result<(), MLError> {
    // Initialize logging
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    println!("=== FactoredQNetwork Q-Value Uniqueness Diagnostic ===\n");

    let device = Device::cuda_if_available(0)?;
    println!("Device: {:?}\n", device);

    // Create factored Q-network
    let network = FactoredQNetwork::new(128, &device)?;

    // Test 1: Single state forward pass
    println!("--- Test 1: Single State Forward Pass ---");
    let state = Tensor::randn(0.0f32, 1.0f32, (1, 128), &device)?;

    let (q_exp, q_ord, q_urg) = network.forward(&state)?;

    // Extract Q-values from each head
    let exp_vec = q_exp.flatten_all()?.to_vec1::<f32>()?;
    let ord_vec = q_ord.flatten_all()?.to_vec1::<f32>()?;
    let urg_vec = q_urg.flatten_all()?.to_vec1::<f32>()?;

    println!("Exposure Q-values (5): {:?}", exp_vec);
    println!("Order Q-values (3): {:?}", ord_vec);
    println!("Urgency Q-values (3): {:?}", urg_vec);

    // Count unique values per head
    let exp_unique: HashSet<_> = exp_vec.iter().map(|&x| (x * 1000.0) as i64).collect();
    let ord_unique: HashSet<_> = ord_vec.iter().map(|&x| (x * 1000.0) as i64).collect();
    let urg_unique: HashSet<_> = urg_vec.iter().map(|&x| (x * 1000.0) as i64).collect();

    println!("\nUnique values per head:");
    println!("  Exposure: {}/5", exp_unique.len());
    println!("  Order: {}/3", ord_unique.len());
    println!("  Urgency: {}/3", urg_unique.len());

    // Test 2: Compute joint Q-values (additive factorization)
    println!("\n--- Test 2: Additive Factorization (45 Joint Q-Values) ---");
    let joint_q = network.compute_joint_q(&q_exp, &q_ord, &q_urg)?;
    let joint_vec = joint_q.flatten_all()?.to_vec1::<f32>()?;

    println!("Joint Q-values shape: {:?}", joint_q.dims());
    println!(
        "Joint Q-values (first 10): {:?}",
        &joint_vec[..10.min(joint_vec.len())]
    );
    println!(
        "Joint Q-values (last 10): {:?}",
        &joint_vec[joint_vec.len().saturating_sub(10)..]
    );

    // Statistics
    let min_q = joint_vec.iter().cloned().fold(f32::INFINITY, f32::min);
    let max_q = joint_vec.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let mean_q = joint_vec.iter().sum::<f32>() / joint_vec.len() as f32;
    let variance: f32 = joint_vec
        .iter()
        .map(|&q| {
            let diff = q - mean_q;
            diff * diff
        })
        .sum::<f32>()
        / joint_vec.len() as f32;
    let std_dev = variance.sqrt();

    println!("\nJoint Q-value Statistics:");
    println!("  Min: {:.4}", min_q);
    println!("  Max: {:.4}", max_q);
    println!("  Range: {:.4}", max_q - min_q);
    println!("  Mean: {:.4}", mean_q);
    println!("  Std Dev: {:.4}", std_dev);

    // Test 3: Uniqueness analysis
    println!("\n--- Test 3: Uniqueness Analysis ---");

    // Count unique values (tolerance: 0.001)
    let unique_values: HashSet<_> = joint_vec.iter().map(|&x| (x * 1000.0) as i64).collect();

    println!("Unique joint Q-values: {}/45", unique_values.len());

    if unique_values.len() < 45 {
        println!(
            "⚠️  WARNING: Only {} unique values detected (expected 45)",
            unique_values.len()
        );
        println!("    This indicates value repetition in the action space.");
    } else {
        println!("✅ All 45 Q-values are unique (within tolerance)");
    }

    // Test 4: Manual factorization verification
    println!("\n--- Test 4: Manual Factorization Verification ---");
    println!("Verifying: Q(exp, ord, urg) = Q_exp[e] + Q_ord[o] + Q_urg[u]");

    // Manually compute first 10 joint Q-values
    let mut manual_q = Vec::new();
    for exp_idx in 0..5 {
        for ord_idx in 0..3 {
            for urg_idx in 0..3 {
                let q_value = exp_vec[exp_idx] + ord_vec[ord_idx] + urg_vec[urg_idx];
                manual_q.push(q_value);

                if manual_q.len() <= 10 {
                    let joint_idx = exp_idx * 9 + ord_idx * 3 + urg_idx;
                    let expected = joint_vec[joint_idx];
                    let diff = (q_value - expected).abs();
                    println!(
                        "  Action[{}] = exp[{}] + ord[{}] + urg[{}] = {:.4} + {:.4} + {:.4} = {:.4} (expected: {:.4}, diff: {:.6})",
                        joint_idx, exp_idx, ord_idx, urg_idx,
                        exp_vec[exp_idx], ord_vec[ord_idx], urg_vec[urg_idx],
                        q_value, expected, diff
                    );
                }
            }
        }
    }

    // Compare manual vs network computation
    let max_diff = manual_q
        .iter()
        .zip(joint_vec.iter())
        .map(|(&manual, &network)| (manual - network).abs())
        .fold(0.0f32, f32::max);

    println!("\nMax difference (manual vs network): {:.6}", max_diff);

    if max_diff < 1e-5 {
        println!("✅ Additive factorization formula is correct");
    } else {
        println!("⚠️  WARNING: Factorization formula mismatch (diff > 1e-5)");
    }

    // Test 5: Value distribution analysis
    println!("\n--- Test 5: Value Distribution Analysis ---");

    // Count how many times each unique value appears
    let mut value_counts: std::collections::HashMap<i64, usize> = std::collections::HashMap::new();
    for &val in &joint_vec {
        let key = (val * 1000.0) as i64;
        *value_counts.entry(key).or_insert(0) += 1;
    }

    // Find duplicate values
    let mut duplicates: Vec<_> = value_counts
        .iter()
        .filter(|(_, &count)| count > 1)
        .collect();
    duplicates.sort_by_key(|(_, &count)| std::cmp::Reverse(count));

    if !duplicates.is_empty() {
        println!("Duplicate Q-values detected:");
        for (val, count) in duplicates.iter().take(5) {
            println!(
                "  Value {:.3} appears {} times",
                (**val as f32) / 1000.0,
                count
            );
        }
    } else {
        println!("✅ No duplicate Q-values detected");
    }

    // Test 6: Batch consistency
    println!("\n--- Test 6: Batch Consistency (32 identical states) ---");

    let batch_state = state.repeat((32, 1))?;
    let (q_exp_batch, q_ord_batch, q_urg_batch) = network.forward(&batch_state)?;
    let joint_q_batch = network.compute_joint_q(&q_exp_batch, &q_ord_batch, &q_urg_batch)?;

    // Check if all batch items have same Q-values
    let batch_vec = joint_q_batch.flatten_all()?.to_vec1::<f32>()?;
    let batch_unique_per_action = (0..45)
        .map(|action_idx| {
            let values: HashSet<_> = (0..32)
                .map(|batch_idx| {
                    let idx = batch_idx * 45 + action_idx;
                    (batch_vec[idx] * 1000.0) as i64
                })
                .collect();
            values.len()
        })
        .collect::<Vec<_>>();

    let all_consistent = batch_unique_per_action.iter().all(|&count| count == 1);

    if all_consistent {
        println!("✅ Batch consistency verified (all 32 items have identical Q-values)");
    } else {
        println!("⚠️  WARNING: Batch inconsistency detected");
        for (action_idx, &unique_count) in batch_unique_per_action.iter().enumerate() {
            if unique_count > 1 {
                println!(
                    "  Action[{}] has {} unique values across batch",
                    action_idx, unique_count
                );
            }
        }
    }

    // Final Summary
    println!("\n=== DIAGNOSTIC SUMMARY ===");
    println!("✅ Network output shapes correct: [1, 5], [1, 3], [1, 3]");
    println!("✅ Joint Q-values shape: [1, 45]");

    if unique_values.len() == 45 {
        println!("✅ All 45 Q-values are unique");
    } else {
        println!("❌ Only {}/45 Q-values are unique", unique_values.len());
    }

    if max_diff < 1e-5 {
        println!("✅ Additive factorization formula verified");
    } else {
        println!(
            "❌ Factorization formula has errors (max diff: {:.6})",
            max_diff
        );
    }

    if all_consistent {
        println!("✅ Batch processing is consistent");
    } else {
        println!("❌ Batch processing has inconsistencies");
    }

    Ok(())
}