//! Test suite for GRN weight initialization verification
//!
//! This test verifies that Gated Residual Network (GRN) layers use proper
//! Xavier/Kaiming weight initialization instead of zeros.
//!
//! Context: Wave 8.6 - Verify that candle_nn::linear() properly initializes
//! weights following Xavier Uniform distribution by default.
//!
//! CRITICAL: Use VarBuilder::from_varmap() for proper weight initialization,
//! NOT VarBuilder::zeros() which creates all-zero weights.

use candle_core::{DType, Device, Tensor};
use candle_nn::{VarBuilder, VarMap};
use std::sync::Arc;

use ml::tft::gated_residual::{GRNStack, GatedLinearUnit, GatedResidualNetwork};
use ml::MLError;

/// Calculate mean of a tensor
fn calculate_mean(tensor: &Tensor) -> Result<f32, MLError> {
    let flat = tensor.flatten_all()?;
    let vec = flat.to_vec1::<f32>()?;
    Ok(vec.iter().sum::<f32>() / vec.len() as f32)
}

/// Calculate standard deviation of a tensor
fn calculate_std_dev(tensor: &Tensor) -> Result<f32, MLError> {
    let flat = tensor.flatten_all()?;
    let vec = flat.to_vec1::<f32>()?;
    let mean = vec.iter().sum::<f32>() / vec.len() as f32;
    let variance = vec.iter().map(|&x| (x - mean).powi(2)).sum::<f32>() / vec.len() as f32;
    Ok(variance.sqrt())
}

/// Calculate min and max values
fn calculate_range(tensor: &Tensor) -> Result<(f32, f32), MLError> {
    let flat = tensor.flatten_all()?;
    let vec = flat.to_vec1::<f32>()?;
    let min = vec.iter().copied().fold(f32::INFINITY, f32::min);
    let max = vec.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    Ok((min, max))
}

#[test]
fn test_grn_weight_initialization_statistics() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let grn = GatedResidualNetwork::new(64, 64, vs.pp("test"))?;

    // Create test input to extract weight information
    let input_data = vec![1.0f32; 128]; // 2 * 64
    let inputs = Tensor::from_slice(&input_data, (2, 64), &device)?;

    // Forward pass to ensure weights are initialized
    let output = grn.forward(&inputs, None)?;

    // Check output statistics
    let mean = calculate_mean(&output)?;
    let std_dev = calculate_std_dev(&output)?;
    let (min, max) = calculate_range(&output)?;

    println!("GRN Output Statistics:");
    println!("  Mean: {:.6}", mean);
    println!("  Std Dev: {:.6}", std_dev);
    println!("  Range: [{:.6}, {:.6}]", min, max);

    // Verify non-zero outputs (would be zero if weights were not initialized)
    assert!(
        std_dev > 0.01,
        "Output std dev should be non-zero (got {}), indicating proper weight initialization",
        std_dev
    );

    // Verify output has reasonable range (not all zeros or infinities)
    assert!(
        min.is_finite() && max.is_finite(),
        "Output should be finite"
    );
    assert!(
        (max - min) > 0.1,
        "Output should have non-trivial range (got {})",
        max - min
    );

    Ok(())
}

#[test]
fn test_grn_different_dims_weight_initialization() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Test with different input/output dimensions (triggers skip_projection)
    let grn = GatedResidualNetwork::new(128, 64, vs.pp("test"))?;

    let input_data = vec![1.0f32; 225]; // 2 * 128
    let inputs = Tensor::from_slice(&input_data, (2, 128), &device)?;

    let output = grn.forward(&inputs, None)?;

    // Check output statistics
    let mean = calculate_mean(&output)?;
    let std_dev = calculate_std_dev(&output)?;

    println!("GRN (different dims) Output Statistics:");
    println!("  Mean: {:.6}", mean);
    println!("  Std Dev: {:.6}", std_dev);

    // Verify skip projection is also properly initialized
    assert!(
        std_dev > 0.01,
        "Output std dev should be non-zero with skip projection (got {})",
        std_dev
    );

    Ok(())
}

#[test]
fn test_grn_context_projection_initialization() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let grn = GatedResidualNetwork::new(64, 64, vs.pp("test"))?;

    let input_data = vec![1.0f32; 128]; // 2 * 64
    let inputs = Tensor::from_slice(&input_data, (2, 64), &device)?;

    let context_data = vec![0.5f32; 128]; // 2 * 64
    let context = Tensor::from_slice(&context_data, (2, 64), &device)?;

    // Forward pass with context
    let output_with_context = grn.forward(&inputs, Some(&context))?;

    // Forward pass without context
    let output_no_context = grn.forward(&inputs, None)?;

    // Check that context has an effect (would be same if context_projection not initialized)
    let diff = (output_with_context - output_no_context)?;
    let diff_std = calculate_std_dev(&diff)?;

    println!("Context effect std dev: {:.6}", diff_std);

    assert!(
        diff_std > 0.01,
        "Context should have measurable effect (got std dev {}), indicating context_projection is initialized",
        diff_std
    );

    Ok(())
}

#[test]
fn test_glu_weight_initialization() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let glu = GatedLinearUnit::new(64, 32, vs.pp("test"))?;

    let input_data = vec![1.0f32; 128]; // 2 * 64
    let inputs = Tensor::from_slice(&input_data, (2, 64), &device)?;

    let output = glu.forward(&inputs)?;

    // Check output statistics
    let mean = calculate_mean(&output)?;
    let std_dev = calculate_std_dev(&output)?;

    println!("GLU Output Statistics:");
    println!("  Mean: {:.6}", mean);
    println!("  Std Dev: {:.6}", std_dev);

    // GLU uses sigmoid gating, so outputs should be in reasonable range
    assert!(
        std_dev > 0.01,
        "GLU output should have non-zero variance (got {})",
        std_dev
    );

    // Check that output is bounded (sigmoid gate keeps values reasonable)
    let (min, max) = calculate_range(&output)?;
    println!("  Range: [{:.6}, {:.6}]", min, max);
    assert!(
        min.is_finite() && max.is_finite(),
        "GLU output should be finite"
    );

    Ok(())
}

#[test]
fn test_grn_stack_weight_initialization() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let stack = GRNStack::new(64, 32, 16, 3, vs.pp("test"))?;

    let input_data = vec![1.0f32; 128]; // 2 * 64
    let inputs = Tensor::from_slice(&input_data, (2, 64), &device)?;

    let output = stack.forward(&inputs, None)?;

    // Check final output statistics
    let mean = calculate_mean(&output)?;
    let std_dev = calculate_std_dev(&output)?;

    println!("GRN Stack Output Statistics:");
    println!("  Mean: {:.6}", mean);
    println!("  Std Dev: {:.6}", std_dev);

    // Multi-layer stack should still have non-zero, finite outputs
    assert!(
        std_dev > 0.01,
        "GRN stack output should have non-zero variance (got {})",
        std_dev
    );

    let (min, max) = calculate_range(&output)?;
    assert!(
        min.is_finite() && max.is_finite(),
        "GRN stack output should be finite"
    );

    Ok(())
}

#[test]
fn test_grn_multiple_forward_passes() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let grn = GatedResidualNetwork::new(32, 32, vs.pp("test"))?;

    // Multiple forward passes with different inputs should produce different outputs
    let input1_data = vec![1.0f32; 64]; // 2 * 32
    let input1 = Tensor::from_slice(&input1_data, (2, 32), &device)?;

    let input2_data = vec![2.0f32; 64]; // 2 * 32
    let input2 = Tensor::from_slice(&input2_data, (2, 32), &device)?;

    let output1 = grn.forward(&input1, None)?;
    let output2 = grn.forward(&input2, None)?;

    // Outputs should be different for different inputs
    let diff = (output2 - output1)?;
    let diff_std = calculate_std_dev(&diff)?;

    println!("Output difference std dev: {:.6}", diff_std);

    assert!(
        diff_std > 0.1,
        "Different inputs should produce different outputs (got std dev {})",
        diff_std
    );

    Ok(())
}

#[test]
fn test_grn_3d_tensor_weight_initialization() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let grn = GatedResidualNetwork::new(16, 16, vs.pp("test"))?;

    // 3D input: [batch_size=2, seq_len=5, hidden_dim=16]
    let input_data = vec![1.0f32; 160]; // 2 * 5 * 16
    let inputs = Tensor::from_slice(&input_data, (2, 5, 16), &device)?;

    let output = grn.forward(&inputs, None)?;

    // Check statistics across all dimensions
    let mean = calculate_mean(&output)?;
    let std_dev = calculate_std_dev(&output)?;

    println!("GRN 3D Output Statistics:");
    println!("  Mean: {:.6}", mean);
    println!("  Std Dev: {:.6}", std_dev);

    assert!(
        std_dev > 0.01,
        "3D tensor output should have non-zero variance (got {})",
        std_dev
    );

    Ok(())
}

#[test]
fn test_grn_batch_consistency() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let grn = GatedResidualNetwork::new(32, 32, vs.pp("test"))?;

    // Create two identical samples in a batch
    let mut input_data = vec![1.0f32; 64]; // 2 * 32
                                           // Make second sample different
    for i in 32..64 {
        input_data[i] = 2.0;
    }
    let inputs = Tensor::from_slice(&input_data, (2, 32), &device)?;

    let output = grn.forward(&inputs, None)?;

    // Extract individual batch elements
    let output_vec = output.to_vec2::<f32>()?;
    let sample1 = &output_vec[0];
    let sample2 = &output_vec[1];

    // Calculate difference between samples
    let diff: Vec<f32> = sample1
        .iter()
        .zip(sample2.iter())
        .map(|(a, b)| (a - b).abs())
        .collect();
    let diff_mean = diff.iter().sum::<f32>() / diff.len() as f32;

    println!("Batch sample difference mean: {:.6}", diff_mean);

    // Different inputs should produce different outputs
    assert!(
        diff_mean > 0.01,
        "Different batch samples should produce different outputs (got mean diff {})",
        diff_mean
    );

    Ok(())
}

#[test]
fn test_grn_zero_input_response() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let grn = GatedResidualNetwork::new(32, 32, vs.pp("test"))?;

    // Zero input
    let zero_input = Tensor::zeros((2, 32), DType::F32, &device)?;
    let output = grn.forward(&zero_input, None)?;

    // Output should not be all zeros if weights are initialized
    // (bias terms and residual connection should produce non-zero output)
    let std_dev = calculate_std_dev(&output)?;

    println!("Zero input output std dev: {:.6}", std_dev);

    // Note: Even with zero input, properly initialized network should have
    // some non-zero response due to bias terms and layer normalization
    let (min, max) = calculate_range(&output)?;
    assert!(
        min.is_finite() && max.is_finite(),
        "Output should be finite even with zero input"
    );

    Ok(())
}