//! Diagnostic test for new direct 45-output FactoredQNetwork architecture
//!
//! Validates that the network produces 45 unique Q-values (not 8 clustered values).

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

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== FactoredQNetwork Architecture Validation ===\n");

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

    // Initialize network
    let network = FactoredQNetwork::new(128, &device)?;
    println!("Network initialized successfully");
    println!("  - State dimension: 128");
    println!("  - Hidden dimension: {}", network.hidden_dim());
    println!("  - Output dimension: 45\n");

    // Generate random state
    let state = Tensor::randn(0.0f32, 1.0f32, (1, 128), &device)?;
    println!("Generated random state with shape: {:?}\n", state.dims());

    // Run forward pass 10 times to check Q-value diversity
    println!("Running 10 forward passes to check Q-value diversity...\n");

    let mut all_unique_counts = Vec::new();

    for i in 0..10 {
        let q_values = network.forward(&state)?;

        // Extract Q-values to vector
        let q_vec = q_values.flatten_all()?.to_vec1::<f32>()?;

        // Count unique Q-values (with 1e-6 tolerance for floating point comparison)
        let mut unique_values = HashSet::new();
        for &q in &q_vec {
            let rounded = (q * 1e6).round() as i64;
            unique_values.insert(rounded);
        }

        let unique_count = unique_values.len();
        all_unique_counts.push(unique_count);

        println!(
            "  Pass {}: {} unique Q-values out of 45",
            i + 1,
            unique_count
        );

        // Print first 10 Q-values for inspection
        print!("    First 10 Q-values: [");
        for (j, &q) in q_vec.iter().take(10).enumerate() {
            if j > 0 {
                print!(", ");
            }
            print!("{:.4}", q);
        }
        println!("]");
    }

    println!();

    // Compute statistics
    let avg_unique: f64 =
        all_unique_counts.iter().sum::<usize>() as f64 / all_unique_counts.len() as f64;
    let min_unique = *all_unique_counts.iter().min().unwrap();
    let max_unique = *all_unique_counts.iter().max().unwrap();

    println!("=== Q-Value Diversity Statistics ===");
    println!("  Average unique Q-values: {:.1}", avg_unique);
    println!("  Minimum unique Q-values: {}", min_unique);
    println!("  Maximum unique Q-values: {}", max_unique);
    println!();

    // Validation
    if avg_unique >= 40.0 {
        println!(
            "✅ SUCCESS: {} unique Q-values confirmed ({:.1}% diversity)",
            avg_unique,
            (avg_unique / 45.0) * 100.0
        );
        println!("   Network architecture is working correctly!");
        println!("   Expected: 45 unique values");
        println!("   Actual:   {:.1} average unique values", avg_unique);
        println!();
        println!("   This confirms the direct 45-output architecture prevents");
        println!("   the additive factorization clustering bug (8 values).");
    } else {
        println!(
            "❌ FAILURE: Only {} unique Q-values detected ({:.1}% diversity)",
            avg_unique,
            (avg_unique / 45.0) * 100.0
        );
        println!("   Network may still have clustering issues!");
        println!("   Expected: >= 40 unique values");
        println!("   Actual:   {:.1} average unique values", avg_unique);
        println!();
        println!("   Action required: Investigate network initialization or forward pass.");
    }

    Ok(())
}