//! Diagnostic test for FactoredQNetwork Q-value diversity
//!
//! Analyzes whether additive factorization Q(e,o,u) = Q_e[e] + Q_o[o] + Q_u[u]
//! produces 45 unique Q-values or creates duplicates/clustering.
//!
//! Expected results:
//! - If 45 unique values: Issue is in argmax/selection logic
//! - If <45 unique values: Additive factorization creates duplicates
//!
//! Run with:
//! ```bash
//! cargo run -p ml --example test_factored_q_values --release
//! ```

use candle_core::{Device, Tensor};
use std::collections::{HashMap, HashSet};

// Inline minimal FactoredQNetwork for standalone diagnostic
use candle_nn::{Linear, Module, VarBuilder, VarMap};

struct SimpleFactoredQNetwork {
    shared_encoder: Linear,
    exposure_head: Linear,
    order_head: Linear,
    urgency_head: Linear,
    device: Device,
}

impl SimpleFactoredQNetwork {
    fn new(state_dim: usize, device: &Device) -> Result<Self, Box<dyn std::error::Error>> {
        let varmap = VarMap::new();
        let vb = VarBuilder::from_varmap(&varmap, candle_core::DType::F32, device);

        let shared_encoder = candle_nn::linear(state_dim, 64, vb.pp("shared_encoder"))?;
        let exposure_head = candle_nn::linear(64, 5, vb.pp("exposure_head"))?;
        let order_head = candle_nn::linear(64, 3, vb.pp("order_head"))?;
        let urgency_head = candle_nn::linear(64, 3, vb.pp("urgency_head"))?;

        Ok(Self {
            shared_encoder,
            exposure_head,
            order_head,
            urgency_head,
            device: device.clone(),
        })
    }

    fn forward(
        &self,
        state: &Tensor,
    ) -> Result<(Tensor, Tensor, Tensor), Box<dyn std::error::Error>> {
        let hidden = self.shared_encoder.forward(state)?;
        let hidden = hidden.relu()?;

        let q_exposure = self.exposure_head.forward(&hidden)?;
        let q_order = self.order_head.forward(&hidden)?;
        let q_urgency = self.urgency_head.forward(&hidden)?;

        Ok((q_exposure, q_order, q_urgency))
    }

    fn compute_joint_q(
        &self,
        q_exposure: &Tensor,
        q_order: &Tensor,
        q_urgency: &Tensor,
    ) -> Result<Tensor, Box<dyn std::error::Error>> {
        let batch_size = q_exposure.dim(0)?;

        // Reshape to [batch, 5, 1, 1]
        let q_exp = q_exposure.reshape((batch_size, 5, 1, 1))?;
        // Reshape to [batch, 1, 3, 1]
        let q_ord = q_order.reshape((batch_size, 1, 3, 1))?;
        // Reshape to [batch, 1, 1, 3]
        let q_urg = q_urgency.reshape((batch_size, 1, 1, 3))?;

        // Broadcast and sum: [batch, 5, 3, 3]
        let joint_q = q_exp.broadcast_add(&q_ord)?;
        let joint_q = joint_q.broadcast_add(&q_urg)?;

        // Flatten to [batch, 45]
        let joint_q = joint_q.reshape((batch_size, 45))?;

        Ok(joint_q)
    }
}

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

    // Use CPU for reproducibility
    let device = Device::Cpu;
    println!("Device: CPU (for reproducibility)");

    // Create network
    let network = SimpleFactoredQNetwork::new(128, &device)?;
    println!("Network created: 128 input → 64 hidden → [5, 3, 3] heads\n");

    // Test 1: Single random state
    println!("=== Test 1: Single Random State ===");
    let state = Tensor::randn(0.0f32, 1.0f32, (1, 128), &device)?;

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

    // Extract raw Q-values from each head
    let exp_values = q_exp.to_vec2::<f32>()?[0].clone();
    let ord_values = q_ord.to_vec2::<f32>()?[0].clone();
    let urg_values = q_urg.to_vec2::<f32>()?[0].clone();

    println!("Exposure Q-values (5): {:?}", exp_values);
    println!("Order Q-values (3):    {:?}", ord_values);
    println!("Urgency Q-values (3):  {:?}", urg_values);

    // Compute joint Q-values using additive factorization
    let joint_q = network.compute_joint_q(&q_exp, &q_ord, &q_urg)?;
    let joint_values = joint_q.to_vec2::<f32>()?[0].clone();

    println!("\nJoint Q-values (45):");
    for (i, val) in joint_values.iter().enumerate() {
        if i % 9 == 0 {
            println!();
        }
        print!("{:8.4} ", val);
    }
    println!("\n");

    // Analyze uniqueness (with epsilon tolerance for floating-point)
    let epsilon = 1e-6;
    let mut unique_values = HashSet::new();
    let mut value_counts = HashMap::new();

    for &val in &joint_values {
        // Round to 6 decimal places for uniqueness check
        let rounded = (val / epsilon).round() as i64;
        unique_values.insert(rounded);
        *value_counts.entry(rounded).or_insert(0) += 1;
    }

    println!("Unique Q-values: {}/45", unique_values.len());
    println!(
        "Duplicate groups: {}",
        value_counts.iter().filter(|(_, &count)| count > 1).count()
    );

    // Show distribution
    let mut sorted_counts: Vec<_> = value_counts.iter().collect();
    sorted_counts.sort_by_key(|(val, _)| *val);

    println!("\nQ-value distribution (rounded to 6 decimals):");
    for (val, count) in sorted_counts.iter().take(10) {
        let actual_val = (*val as f32) * epsilon;
        println!("  Q={:8.4}: appears {} times", actual_val, count);
    }
    if sorted_counts.len() > 10 {
        println!("  ... ({} more unique values)", sorted_counts.len() - 10);
    }

    // Test 2: Verify additive factorization formula
    println!("\n=== Test 2: Manual Verification of Additive Formula ===");

    // Manually compute first 5 joint Q-values and compare with network output
    println!("Verifying Q(e,o,u) = Q_e[e] + Q_o[o] + Q_u[u]:");

    for idx in 0..5 {
        let exp_idx = idx / 9;
        let ord_idx = (idx % 9) / 3;
        let urg_idx = idx % 3;

        let manual_q = exp_values[exp_idx] + ord_values[ord_idx] + urg_values[urg_idx];
        let network_q = joint_values[idx];
        let diff = (manual_q - network_q).abs();

        println!(
            "  Index {}: Q_e[{}] + Q_o[{}] + Q_u[{}] = {:.4} + {:.4} + {:.4} = {:.4} (network: {:.4}, diff: {:.6})",
            idx, exp_idx, ord_idx, urg_idx,
            exp_values[exp_idx], ord_values[ord_idx], urg_values[urg_idx],
            manual_q, network_q, diff
        );

        if diff > 1e-5 {
            println!("    WARNING: Mismatch detected!");
        }
    }

    // Test 3: Multiple random initializations
    println!("\n=== Test 3: Average Uniqueness Across 100 Random States ===");

    let mut total_unique = 0;
    let mut min_unique = 45;
    let mut max_unique = 0;

    for trial in 0..100 {
        let state = Tensor::randn(0.0f32, 1.0f32, (1, 128), &device)?;
        let (q_exp, q_ord, q_urg) = network.forward(&state)?;
        let joint_q = network.compute_joint_q(&q_exp, &q_ord, &q_urg)?;
        let joint_values = joint_q.to_vec2::<f32>()?[0].clone();

        let mut unique_values = HashSet::new();
        for &val in &joint_values {
            let rounded = (val / epsilon).round() as i64;
            unique_values.insert(rounded);
        }

        let unique_count = unique_values.len();
        total_unique += unique_count;
        min_unique = min_unique.min(unique_count);
        max_unique = max_unique.max(unique_count);

        if trial < 10 {
            println!("  Trial {}: {}/45 unique values", trial, unique_count);
        }
    }

    let avg_unique = total_unique as f32 / 100.0;
    println!("\nStatistics over 100 trials:");
    println!("  Average unique values: {:.2}/45", avg_unique);
    println!("  Min unique values: {}/45", min_unique);
    println!("  Max unique values: {}/45", max_unique);

    // Test 4: Analyze theoretical worst case
    println!("\n=== Test 4: Theoretical Analysis ===");
    println!("Additive factorization: Q(e,o,u) = Q_e[e] + Q_o[o] + Q_u[u]");
    println!("Number of possible sums: 5 (exposure) × 3 (order) × 3 (urgency) = 45");
    println!("\nHowever, if the head outputs are similar in magnitude (e.g., all near 0.0),");
    println!("many combinations can produce identical or very close sums due to:");
    println!("  1. Limited precision (floating-point rounding)");
    println!("  2. Similar weight initialization (default init)");
    println!("  3. Small variance in early training");

    // Simulate worst case: all heads output near-zero
    println!("\nSimulating worst case (all heads output ~0.0):");
    let zero_exp = vec![0.0f32, 0.01, 0.02, 0.03, 0.04];
    let zero_ord = vec![0.00f32, 0.01, 0.02];
    let zero_urg = vec![0.00f32, 0.01, 0.02];

    let mut worst_case_unique = HashSet::new();
    for exp in &zero_exp {
        for ord in &zero_ord {
            for urg in &zero_urg {
                let sum = exp + ord + urg;
                let rounded = (sum / epsilon).round() as i64;
                worst_case_unique.insert(rounded);
            }
        }
    }

    println!(
        "  Unique sums in worst case: {}/45",
        worst_case_unique.len()
    );

    // Test 5: Recommendation
    println!("\n=== Diagnosis Summary ===");

    if avg_unique < 20.0 {
        println!(
            "❌ CRITICAL: Additive factorization produces severe clustering (<20 unique values)"
        );
        println!("\nRecommended fixes:");
        println!("  1. Multiplicative factorization: Q(e,o,u) = Q_e[e] × Q_o[o] × Q_u[u]");
        println!("     Pros: More expressive, less clustering");
        println!("     Cons: Requires Q-values > 0, harder to train");
        println!("\n  2. Concatenation + single head: [hidden, 64] → [45] directly");
        println!("     Pros: Full expressiveness, guaranteed 45 unique values");
        println!("     Cons: Loss of factored structure, no sub-action interpretability");
        println!("\n  3. Weighted sum with learnable weights: Q = w1*Q_e + w2*Q_o + w3*Q_u");
        println!("     Pros: Retains additive structure, learnable importance");
        println!("     Cons: Still susceptible to clustering if weights are similar");
        println!("\n  4. Increase variance via initialization (σ=0.5 instead of default)");
        println!("     Pros: Simple fix, retains additive structure");
        println!("     Cons: May cause training instability, only delays clustering");
    } else if avg_unique < 40.0 {
        println!("⚠️  WARNING: Additive factorization produces moderate clustering (20-40 unique values)");
        println!("\nRecommended fixes:");
        println!("  1. Increase head variance via custom initialization");
        println!("  2. Consider multiplicative or concatenation approach");
    } else {
        println!("✅ OK: Additive factorization produces good diversity (40+ unique values)");
        println!("\nIf argmax still selects only 8 actions, the issue is in:");
        println!("  1. Epsilon-greedy exploration logic");
        println!("  2. Argmax tie-breaking (Q-values too close together)");
        println!("  3. Position masking (aggressive filtering)");
    }

    println!("\n=== Action Breakdown ===");
    println!("45 actions = 5 exposure × 3 order × 3 urgency");
    println!("\nExposure levels (5):");
    println!("  0: Short100 (-100%)");
    println!("  1: Short50  (-50%)");
    println!("  2: Flat     (0%)");
    println!("  3: Long50   (+50%)");
    println!("  4: Long100  (+100%)");
    println!("\nOrder types (3):");
    println!("  0: Market      (0.15% fee)");
    println!("  1: LimitMaker  (0.05% fee)");
    println!("  2: IoC         (0.10% fee)");
    println!("\nUrgency levels (3):");
    println!("  0: Patient    (0.5x weight)");
    println!("  1: Normal     (1.0x weight)");
    println!("  2: Aggressive (1.5x weight)");

    Ok(())
}