//! Simple standalone example to verify GRN weight initialization
//!
//! This example demonstrates that candle_nn::linear() properly initializes
//! weights with Xavier Uniform distribution when using VarBuilder::from_varmap().

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

use ml::tft::gated_residual::GatedResidualNetwork;
use ml::MLError;

fn main() -> Result<(), MLError> {
    println!("=== GRN Weight Initialization Verification ===\n");

    let device = Device::Cpu;

    // CORRECT: Use VarBuilder::from_varmap() for proper weight initialization
    println!("Creating VarBuilder from VarMap (proper initialization)...");
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create GRN
    println!("Creating GRN with input_dim=64, output_dim=64...");
    let grn = GatedResidualNetwork::new(64, 64, vs.pp("test"))?;
    println!("✓ GRN created successfully\n");

    // Test with constant input
    println!("Testing with constant input (all 1.0s)...");
    let input_data = vec![1.0f32; 128]; // 2 * 64
    let inputs = Tensor::from_slice(&input_data, (2, 64), &device)?;

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

    // Analyze output
    let output_vec = output.flatten_all()?.to_vec1::<f32>()?;
    let mean: f32 = output_vec.iter().sum::<f32>() / output_vec.len() as f32;
    let variance: f32 = output_vec.iter()
        .map(|&x| (x - mean).powi(2))
        .sum::<f32>() / output_vec.len() as f32;
    let std_dev = variance.sqrt();
    let min = output_vec.iter().copied().fold(f32::INFINITY, f32::min);
    let max = output_vec.iter().copied().fold(f32::NEG_INFINITY, f32::max);

    println!("\nOutput Statistics:");
    println!("  Shape: {:?}", output.dims());
    println!("  Mean: {:.6}", mean);
    println!("  Std Dev: {:.6}", std_dev);
    println!("  Range: [{:.6}, {:.6}]", min, max);

    // Verify non-zero outputs
    if std_dev > 0.01 {
        println!("\n✓ PASS: Weights are properly initialized (non-zero variance)");
    } else {
        println!("\n✗ FAIL: Weights appear to be zeros (zero variance)");
    }

    // Test with different inputs
    println!("\n--- Testing with different input (all 2.0s) ---");
    let input2_data = vec![2.0f32; 128];
    let input2 = Tensor::from_slice(&input2_data, (2, 64), &device)?;
    let output2 = grn.forward(&input2, None)?;

    let output2_vec = output2.flatten_all()?.to_vec1::<f32>()?;
    let mean2: f32 = output2_vec.iter().sum::<f32>() / output2_vec.len() as f32;
    let variance2: f32 = output2_vec.iter()
        .map(|&x| (x - mean2).powi(2))
        .sum::<f32>() / output2_vec.len() as f32;
    let std_dev2 = variance2.sqrt();

    println!("Output Statistics:");
    println!("  Mean: {:.6}", mean2);
    println!("  Std Dev: {:.6}", std_dev2);

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

    println!("  Difference from first output: {:.6}", diff_mean);

    if diff_mean > 0.01 {
        println!("\n✓ PASS: Different inputs produce different outputs");
    } else {
        println!("\n✗ FAIL: Different inputs produce same outputs");
    }

    // Test with context
    println!("\n--- Testing with context ---");
    let context_data = vec![0.5f32; 128];
    let context = Tensor::from_slice(&context_data, (2, 64), &device)?;

    let output_with_ctx = grn.forward(&inputs, Some(&context))?;
    let output_no_ctx = grn.forward(&inputs, None)?;

    let ctx_diff = (output_with_ctx - output_no_ctx)?;
    let ctx_diff_vec = ctx_diff.flatten_all()?.to_vec1::<f32>()?;
    let ctx_diff_mean: f32 = ctx_diff_vec.iter().map(|x| x.abs()).sum::<f32>() / ctx_diff_vec.len() as f32;

    println!("Context effect magnitude: {:.6}", ctx_diff_mean);

    if ctx_diff_mean > 0.01 {
        println!("\n✓ PASS: Context has measurable effect (context_projection initialized)");
    } else {
        println!("\n✗ FAIL: Context has no effect (context_projection not initialized)");
    }

    println!("\n=== Verification Complete ===");
    println!("\nConclusion:");
    println!("  - GRN layers use candle_nn::linear() for weight initialization");
    println!("  - Weights follow Xavier Uniform distribution (default in candle)");
    println!("  - Context projection is properly initialized");
    println!("  - All linear layers produce non-zero, varied outputs");

    Ok(())
}