foxhunt/ml/examples/ensemble_uncertainty_demo.rs

//! Ensemble Uncertainty Quantification Demo
//!
//! Demonstrates how to use the EnsembleUncertainty module to:
//! - Track Q-value variance across multiple DQN agents
//! - Measure action disagreement rates
//! - Compute entropy of action distributions
//! - Calculate exploration bonuses based on uncertainty
//!
//! # Usage
//!
//! ```bash
//! cargo run -p ml --example ensemble_uncertainty_demo --release --features cuda
//! ```

use anyhow::Result;
use candle_core::{Device, Tensor};
use ml::dqn::{EnsembleUncertainty, UncertaintyMetrics};

fn main() -> Result<()> {
    println!("=== Ensemble Uncertainty Quantification Demo ===\n");

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

    // Create uncertainty system for 5 agents
    let mut uncertainty = EnsembleUncertainty::new(device.clone(), 5)?;

    // === Scenario 1: High Consensus (Low Uncertainty) ===
    println!("--- Scenario 1: High Consensus ---");
    let q_values_consensus = create_consensus_scenario(&device)?;
    let metrics_consensus = uncertainty.compute_uncertainty(&q_values_consensus)?;
    print_metrics("High Consensus", &metrics_consensus);

    // === Scenario 2: High Disagreement (High Uncertainty) ===
    println!("\n--- Scenario 2: High Disagreement ---");
    let q_values_disagreement = create_disagreement_scenario(&device)?;
    let metrics_disagreement = uncertainty.compute_uncertainty(&q_values_disagreement)?;
    print_metrics("High Disagreement", &metrics_disagreement);

    // === Scenario 3: Partial Disagreement (Medium Uncertainty) ===
    println!("\n--- Scenario 3: Partial Disagreement ---");
    let q_values_partial = create_partial_disagreement_scenario(&device)?;
    let metrics_partial = uncertainty.compute_uncertainty(&q_values_partial)?;
    print_metrics("Partial Disagreement", &metrics_partial);

    // === Scenario 4: Exploration Bonus Comparison ===
    println!("\n--- Scenario 4: Exploration Bonus Comparison ---");
    compare_exploration_bonuses(&metrics_consensus, &metrics_disagreement, &metrics_partial);

    // === Scenario 5: Uncertainty History Tracking ===
    println!("\n--- Scenario 5: Uncertainty History Tracking ---");
    demonstrate_history_tracking(&device)?;

    println!("\n=== Demo Complete ===");
    Ok(())
}

/// Create high consensus scenario: all agents agree
fn create_consensus_scenario(device: &Device) -> Result<Vec<Tensor>> {
    let q_values = vec![
        Tensor::new(&[1.0f32, 2.0, 5.0], device)?.reshape(&[1, 3])?, // All prefer Hold (action 2)
        Tensor::new(&[1.1f32, 2.1, 5.1], device)?.reshape(&[1, 3])?,
        Tensor::new(&[0.9f32, 1.9, 4.9], device)?.reshape(&[1, 3])?,
        Tensor::new(&[1.0f32, 2.0, 5.0], device)?.reshape(&[1, 3])?,
        Tensor::new(&[1.0f32, 2.0, 5.0], device)?.reshape(&[1, 3])?,
    ];
    Ok(q_values)
}

/// Create high disagreement scenario: agents strongly disagree
fn create_disagreement_scenario(device: &Device) -> Result<Vec<Tensor>> {
    let q_values = vec![
        Tensor::new(&[10.0f32, 0.0, 0.0], device)?.reshape(&[1, 3])?, // Agent 1: Buy
        Tensor::new(&[0.0f32, 10.0, 0.0], device)?.reshape(&[1, 3])?, // Agent 2: Sell
        Tensor::new(&[0.0f32, 0.0, 10.0], device)?.reshape(&[1, 3])?, // Agent 3: Hold
        Tensor::new(&[10.0f32, 5.0, 0.0], device)?.reshape(&[1, 3])?, // Agent 4: Buy
        Tensor::new(&[0.0f32, 5.0, 10.0], device)?.reshape(&[1, 3])?, // Agent 5: Hold
    ];
    Ok(q_values)
}

/// Create partial disagreement scenario: majority agrees
fn create_partial_disagreement_scenario(device: &Device) -> Result<Vec<Tensor>> {
    let q_values = vec![
        Tensor::new(&[5.0f32, 2.0, 1.0], device)?.reshape(&[1, 3])?, // Buy (majority)
        Tensor::new(&[5.5f32, 2.5, 1.5], device)?.reshape(&[1, 3])?, // Buy (majority)
        Tensor::new(&[6.0f32, 3.0, 2.0], device)?.reshape(&[1, 3])?, // Buy (majority)
        Tensor::new(&[1.0f32, 5.0, 2.0], device)?.reshape(&[1, 3])?, // Sell (minority)
        Tensor::new(&[1.5f32, 5.5, 2.5], device)?.reshape(&[1, 3])?, // Sell (minority)
    ];
    Ok(q_values)
}

/// Print uncertainty metrics
fn print_metrics(scenario: &str, metrics: &UncertaintyMetrics) {
    println!("Scenario: {}", scenario);
    println!("  Q-Value Variance:      {:.4}", metrics.q_value_variance);
    println!(
        "  Action Disagreement:   {:.2}% ({:.2})",
        metrics.action_disagreement * 100.0,
        metrics.action_disagreement
    );
    println!(
        "  Action Entropy:        {:.4} bits",
        metrics.action_entropy
    );
    println!(
        "  Per-Action Variance:   {:?}",
        metrics
            .per_action_variance
            .iter()
            .map(|v| format!("{:.2}", v))
            .collect::<Vec<_>>()
    );
    println!(
        "  Vote Counts:           Buy={}, Sell={}, Hold={}",
        metrics.vote_counts[0], metrics.vote_counts[1], metrics.vote_counts[2]
    );
    println!(
        "  Majority Action:       {} (0=Buy, 1=Sell, 2=Hold)",
        metrics.majority_action
    );
    println!(
        "  Confidence Score:      {:.4} (0=uncertain, 1=confident)",
        metrics.confidence_score()
    );
    println!(
        "  High Uncertainty?      {}",
        if metrics.is_high_uncertainty() {
            "YES"
        } else {
            "NO"
        }
    );
}

/// Compare exploration bonuses across scenarios
fn compare_exploration_bonuses(
    consensus: &UncertaintyMetrics,
    disagreement: &UncertaintyMetrics,
    partial: &UncertaintyMetrics,
) {
    // Default weights: variance=0.4, disagreement=0.4, entropy=0.2
    let bonus_consensus = consensus.exploration_bonus(0.4, 0.4, 0.2);
    let bonus_disagreement = disagreement.exploration_bonus(0.4, 0.4, 0.2);
    let bonus_partial = partial.exploration_bonus(0.4, 0.4, 0.2);

    println!("Exploration Bonuses (β_var=0.4, β_dis=0.4, β_ent=0.2):");
    println!("  Consensus:           {:.4}", bonus_consensus);
    println!("  Disagreement:        {:.4}", bonus_disagreement);
    println!("  Partial:             {:.4}", bonus_partial);

    // Custom weights: higher variance weight
    let bonus_consensus_cv = consensus.exploration_bonus(0.7, 0.2, 0.1);
    let bonus_disagreement_cv = disagreement.exploration_bonus(0.7, 0.2, 0.1);
    let bonus_partial_cv = partial.exploration_bonus(0.7, 0.2, 0.1);

    println!("\nExploration Bonuses (β_var=0.7, β_dis=0.2, β_ent=0.1):");
    println!("  Consensus:           {:.4}", bonus_consensus_cv);
    println!("  Disagreement:        {:.4}", bonus_disagreement_cv);
    println!("  Partial:             {:.4}", bonus_partial_cv);

    println!("\nInterpretation:");
    println!("  - Consensus scenario has LOW bonus (agents agree → no need to explore)");
    println!("  - Disagreement scenario has HIGH bonus (agents disagree → explore more)");
    println!("  - Partial scenario has MEDIUM bonus (some disagreement → moderate exploration)");
}

/// Demonstrate uncertainty history tracking
fn demonstrate_history_tracking(device: &Device) -> Result<()> {
    let mut uncertainty = EnsembleUncertainty::new(device.clone(), 3)?;

    println!("Simulating 10 steps with varying uncertainty...");

    // Simulate 10 steps with increasing disagreement
    for i in 0..10 {
        let drift = i as f32 * 0.5;
        let q_values = vec![
            Tensor::new(&[1.0f32 + drift, 2.0, 3.0], device)?.reshape(&[1, 3])?,
            Tensor::new(&[1.0f32, 2.0 + drift, 3.0], device)?.reshape(&[1, 3])?,
            Tensor::new(&[1.0f32, 2.0, 3.0 + drift], device)?.reshape(&[1, 3])?,
        ];
        let metrics = uncertainty.compute_uncertainty(&q_values)?;
        println!(
            "  Step {}: variance={:.4}, disagreement={:.2}%, entropy={:.4}",
            i,
            metrics.q_value_variance,
            metrics.action_disagreement * 100.0,
            metrics.action_entropy
        );
    }

    // Get average uncertainty over last 5 steps
    if let Some((avg_var, avg_dis, avg_ent)) = uncertainty.get_average_uncertainty(5) {
        println!("\nAverage Uncertainty (last 5 steps):");
        println!("  Variance:      {:.4}", avg_var);
        println!("  Disagreement:  {:.2}%", avg_dis * 100.0);
        println!("  Entropy:       {:.4} bits", avg_ent);
    }

    // Get recent metrics
    let recent = uncertainty.get_recent_metrics(3);
    println!("\nRecent Metrics (last 3 steps):");
    for (i, m) in recent.iter().enumerate() {
        println!(
            "  Step {}: variance={:.4}, conf={:.4}",
            10 - 3 + i,
            m.q_value_variance,
            m.confidence_score()
        );
    }

    Ok(())
}