use candle_core::{Device, DType, Tensor};
use candle_nn::{VarBuilder, VarMap, linear};
use ml::tft::{TemporalFusionTransformer, TFTConfig};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== TFT Weight Initialization Checker ===\n");

    // Test 1: Check candle_nn::linear initialization
    println!("Test 1: candle_nn::linear default initialization");
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create a simple linear layer
    let layer = linear(10, 5, vs.pp("test_layer"))?;

    // Get the weight tensor
    let all_tensors: Vec<_> = varmap.all_vars().into_iter().collect();
    for (name, tensor) in all_tensors {
        println!("  Tensor: {}", name);
        println!("  Shape: {:?}", tensor.dims());

        let data = tensor.flatten_all()?.to_vec1::<f32>()?;
        let sum: f32 = data.iter().sum();
        let mean = sum / data.len() as f32;
        let variance: f32 = data.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / data.len() as f32;
        let std_dev = variance.sqrt();

        let all_zeros = data.iter().all(|&x| x.abs() < 1e-10);
        let min = data.iter().cloned().fold(f32::INFINITY, f32::min);
        let max = data.iter().cloned().fold(f32::NEG_INFINITY, f32::max);

        println!("  Mean: {:.6}", mean);
        println!("  Std Dev: {:.6}", std_dev);
        println!("  Min: {:.6}", min);
        println!("  Max: {:.6}", max);
        println!("  All zeros: {}", all_zeros);

        if all_zeros {
            println!("  ❌ WARNING: Weights are ZERO-INITIALIZED");
        } else {
            println!("  ✅ OK: Weights are properly initialized");
        }
        println!();
    }

    // Test 2: Check TFT context enrichment weights
    println!("\nTest 2: TFT static context enrichment weights");
    let config = TFTConfig {
        input_dim: 10,
        hidden_dim: 32,
        num_heads: 2,
        num_layers: 2,
        prediction_horizon: 5,
        sequence_length: 20,
        num_quantiles: 5,
        num_static_features: 3,
        num_known_features: 2,
        num_unknown_features: 5,
        ..Default::default()
    };

    let tft = TemporalFusionTransformer::new(config)?;

    // Create test inputs
    let batch_size = 4;
    let static_features = Tensor::randn(0.0f32, 1.0, (batch_size, 3), &device)?;
    let historical_features = Tensor::randn(0.0f32, 1.0, (batch_size, 20, 5), &device)?;
    let future_features = Tensor::randn(0.0f32, 1.0, (batch_size, 5, 2), &device)?;

    // Run forward pass to see if static context has any effect
    println!("  Running forward pass with static features...");
    let mut tft_with_static = tft;
    let output_with_static = tft_with_static.forward(&static_features, &historical_features, &future_features)?;
    println!("  Output shape: {:?}", output_with_static.dims());

    // Check output values
    let output_data = output_with_static.flatten_all()?.to_vec1::<f32>()?;
    let sum: f32 = output_data.iter().sum();
    let mean = sum / output_data.len() as f32;
    let variance: f32 = output_data.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / output_data.len() as f32;
    let std_dev = variance.sqrt();

    println!("  Output mean: {:.6}", mean);
    println!("  Output std_dev: {:.6}", std_dev);

    let all_zeros = output_data.iter().all(|&x| x.abs() < 1e-10);
    if all_zeros {
        println!("  ❌ CRITICAL: All outputs are ZERO - context enrichment not working!");
    } else {
        println!("  ✅ OK: Outputs are non-zero");
    }

    // Test 3: Compare outputs with different static context
    println!("\nTest 3: Static context effect validation");
    let config2 = TFTConfig {
        input_dim: 10,
        hidden_dim: 32,
        num_heads: 2,
        num_layers: 2,
        prediction_horizon: 5,
        sequence_length: 20,
        num_quantiles: 5,
        num_static_features: 3,
        num_known_features: 2,
        num_unknown_features: 5,
        ..Default::default()
    };

    let tft2 = TemporalFusionTransformer::new(config2)?;

    // Create different static features (all zeros vs all ones)
    let static_zeros = Tensor::zeros((batch_size, 3), DType::F32, &device)?;
    let static_ones = Tensor::ones((batch_size, 3), DType::F32, &device)?;

    let mut tft_test1 = tft2;
    let output_zeros = tft_test1.forward(&static_zeros, &historical_features, &future_features)?;

    let config3 = TFTConfig {
        input_dim: 10,
        hidden_dim: 32,
        num_heads: 2,
        num_layers: 2,
        prediction_horizon: 5,
        sequence_length: 20,
        num_quantiles: 5,
        num_static_features: 3,
        num_known_features: 2,
        num_unknown_features: 5,
        ..Default::default()
    };
    let tft3 = TemporalFusionTransformer::new(config3)?;
    let mut tft_test2 = tft3;
    let output_ones = tft_test2.forward(&static_ones, &historical_features, &future_features)?;

    // Compute difference
    let diff = (&output_ones - &output_zeros)?;
    let diff_data = diff.flatten_all()?.to_vec1::<f32>()?;
    let diff_sum: f32 = diff_data.iter().map(|x| x.abs()).sum();
    let diff_mean = diff_sum / diff_data.len() as f32;

    println!("  Mean absolute difference: {:.6}", diff_mean);

    if diff_mean < 1e-6 {
        println!("  ❌ CRITICAL: Static context has NO effect on predictions!");
        println!("  This indicates zero-initialized or missing context enrichment weights.");
    } else {
        println!("  ✅ OK: Static context affects predictions (difference: {:.6})", diff_mean);
    }

    println!("\n=== Summary ===");
    println!("If weights are zero-initialized, the static context enrichment layer");
    println!("will multiply context features by zero, effectively ignoring them.");
    println!("This matches the observation in the test where static features have no effect.");

    Ok(())
}