//! TFT Static Context Contribution Tests - Wave 8.9
//!
//! Validates that static context features have measurable impact on TFT predictions,
//! following Wave 7.5 analysis findings on architectural imbalance.
//!
//! ## Test Coverage:
//! 1. Basic static context contribution (zeros vs signal)
//! 2. Ablation study (with/without static context)
//! 3. Feature importance (vary individual static features)
//! 4. Context gating validation (non-zero weights)
//! 5. Architectural imbalance analysis (5 static vs 65×241 temporal)
//!
//! ## Expected Results (from Wave 7.5):
//! - Mean absolute difference: 0.001-0.1 (small but non-zero)
//! - Effect is measurable but dominated by temporal features
//! - Context projection weights are non-zero (Xavier initialization)

#![allow(unused_crate_dependencies)]

use candle_core::{DType, Device, Tensor};
use candle_nn::VarBuilder;

use ml::tft::{TemporalFusionTransformer, TFTConfig};
use ml::MLError;

// ============================================================================
// TEST 1: BASIC STATIC CONTEXT CONTRIBUTION
// ============================================================================

#[test]
fn test_tft_static_context_contribution_basic() -> Result<(), MLError> {
    let device = Device::Cpu;

    // Configuration matching production TFT setup
    let config = TFTConfig {
        input_dim: 241,
        hidden_dim: 64,
        num_heads: 4,
        num_layers: 3,
        prediction_horizon: 5,
        sequence_length: 60,
        num_quantiles: 9,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 241,
        dropout_rate: 0.1,
        ..Default::default()
    };

    let mut tft = TemporalFusionTransformer::new(config.clone())?;

    // Create base inputs with realistic dimensions
    let batch_size = 4;
    let historical_data = vec![0.5f32; batch_size * 60 * 241];
    let historical = Tensor::from_slice(&historical_data, (batch_size, 60, 241), &device)?;

    let future_data = vec![0.3f32; batch_size * 5 * 10];
    let future = Tensor::from_slice(&future_data, (batch_size, 5, 10), &device)?;

    // Test 1: Static context all zeros
    let static_zeros = Tensor::zeros((batch_size, 5), DType::F32, &device)?;
    let pred_zeros = tft.forward(&static_zeros, &historical, &future)?;

    // Test 2: Static context with signal (non-zero values)
    let static_signal_data = vec![2.0f32; batch_size * 5];
    let static_signal = Tensor::from_slice(&static_signal_data, (batch_size, 5), &device)?;

    // Need to create new TFT instance for fair comparison (avoid state contamination)
    let mut tft2 = TemporalFusionTransformer::new(config.clone())?;
    let pred_signal = tft2.forward(&static_signal, &historical, &future)?;

    // Compute difference between predictions
    let diff = (pred_signal - pred_zeros)?.abs()?;
    let mean_diff = diff.mean_all()?.to_vec0::<f32>()?;

    println!("Static context mean absolute difference: {:.6}", mean_diff);

    // Static context should have SOME effect (even if small)
    assert!(
        mean_diff > 0.001,
        "Static context should affect predictions (diff: {:.6}, threshold: 0.001)",
        mean_diff
    );

    // Effect should be bounded (not dominating)
    assert!(
        mean_diff < 1.0,
        "Static context effect should be bounded (diff: {:.6}, max: 1.0)",
        mean_diff
    );

    // Additional validation: predictions should have valid shape
    assert_eq!(pred_zeros.dims(), &[batch_size, 5, 9]); // [batch, horizon, quantiles]
    assert_eq!(pred_signal.dims(), &[batch_size, 5, 9]);

    Ok(())
}

// ============================================================================
// TEST 2: ABLATION STUDY - WITH/WITHOUT STATIC CONTEXT
// ============================================================================

#[test]
fn test_tft_static_context_ablation_study() -> Result<(), MLError> {
    let device = Device::Cpu;
    let batch_size = 2;

    // Configuration with static features
    let config_with_static = TFTConfig {
        input_dim: 241,
        hidden_dim: 32,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 3,
        sequence_length: 30,
        num_quantiles: 5,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 241,
        dropout_rate: 0.0, // Disable dropout for reproducibility
        ..Default::default()
    };

    let mut tft_with_static = TemporalFusionTransformer::new(config_with_static.clone())?;

    // Create test inputs
    let historical = Tensor::randn(0.0f32, 1.0, (batch_size, 30, 241), &device)?;
    let future = Tensor::randn(0.0f32, 1.0, (batch_size, 3, 10), &device)?;
    let static_features = Tensor::randn(0.0f32, 0.5, (batch_size, 5), &device)?;

    // Forward pass with static context
    let pred_with_static = tft_with_static.forward(&static_features, &historical, &future)?;

    // Forward pass with null static context (all zeros)
    let static_null = Tensor::zeros((batch_size, 5), DType::F32, &device)?;
    let mut tft_null_static = TemporalFusionTransformer::new(config_with_static.clone())?;
    let pred_null_static = tft_null_static.forward(&static_null, &historical, &future)?;

    // Compute performance difference
    let diff = (pred_with_static - pred_null_static)?.abs()?;
    let mean_diff = diff.mean_all()?.to_vec0::<f32>()?;
    let max_diff = diff.max_keepdim(0)?.max_keepdim(1)?.max_keepdim(2)?.to_vec0::<f32>()?;

    println!("Ablation study results:");
    println!("  Mean difference: {:.6}", mean_diff);
    println!("  Max difference:  {:.6}", max_diff);

    // Static context should contribute measurably
    assert!(
        mean_diff > 0.0001,
        "Static context should have measurable effect (mean diff: {:.6})",
        mean_diff
    );

    // Maximum difference should show localized impact
    assert!(
        max_diff > mean_diff,
        "Max difference ({:.6}) should exceed mean ({:.6})",
        max_diff,
        mean_diff
    );

    // Verify predictions are valid (no NaN/Inf)
    let pred_with_data = pred_with_static.flatten_all()?.to_vec1::<f32>()?;
    let pred_null_data = pred_null_static.flatten_all()?.to_vec1::<f32>()?;

    assert!(pred_with_data.iter().all(|&x| x.is_finite()));
    assert!(pred_null_data.iter().all(|&x| x.is_finite()));

    Ok(())
}

// ============================================================================
// TEST 3: FEATURE IMPORTANCE - INDIVIDUAL STATIC FEATURES
// ============================================================================

#[test]
fn test_tft_static_feature_individual_importance() -> Result<(), MLError> {
    let device = Device::Cpu;
    let batch_size = 2;

    let config = TFTConfig {
        input_dim: 241,
        hidden_dim: 32,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 3,
        sequence_length: 20,
        num_quantiles: 5,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 241,
        dropout_rate: 0.0,
        ..Default::default()
    };

    // Create base inputs
    let historical = Tensor::randn(0.0f32, 1.0, (batch_size, 20, 241), &device)?;
    let future = Tensor::randn(0.0f32, 1.0, (batch_size, 3, 10), &device)?;

    // Baseline: all static features at zero
    let static_baseline = Tensor::zeros((batch_size, 5), DType::F32, &device)?;
    let mut tft_baseline = TemporalFusionTransformer::new(config.clone())?;
    let pred_baseline = tft_baseline.forward(&static_baseline, &historical, &future)?;

    // Test each static feature individually
    let mut feature_impacts = Vec::new();

    for feature_idx in 0..5 {
        // Create static context with only feature_idx set to non-zero
        let mut static_data = vec![0.0f32; batch_size * 5];
        for batch_idx in 0..batch_size {
            static_data[batch_idx * 5 + feature_idx] = 1.0; // Set feature to 1.0
        }
        let static_single = Tensor::from_slice(&static_data, (batch_size, 5), &device)?;

        // Forward pass
        let mut tft_single = TemporalFusionTransformer::new(config.clone())?;
        let pred_single = tft_single.forward(&static_single, &historical, &future)?;

        // Compute impact
        let diff = (pred_single - pred_baseline.clone())?.abs()?;
        let mean_impact = diff.mean_all()?.to_vec0::<f32>()?;

        feature_impacts.push((feature_idx, mean_impact));
        println!("Feature {} impact: {:.6}", feature_idx, mean_impact);
    }

    // Validate that at least some features have measurable impact
    let significant_features = feature_impacts
        .iter()
        .filter(|(_, impact)| *impact > 0.0001)
        .count();

    assert!(
        significant_features > 0,
        "At least one static feature should have measurable impact (found {})",
        significant_features
    );

    // Verify impacts are bounded
    for (idx, impact) in &feature_impacts {
        assert!(
            *impact < 0.5,
            "Feature {} impact should be bounded (got {:.6})",
            idx,
            impact
        );
    }

    Ok(())
}

// ============================================================================
// TEST 4: CONTEXT GATING VALIDATION - NON-ZERO WEIGHTS
// ============================================================================

#[test]
fn test_tft_static_context_projection_active() -> Result<(), MLError> {
    let device = Device::Cpu;

    let config = TFTConfig {
        input_dim: 64,
        hidden_dim: 32,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 3,
        sequence_length: 10,
        num_quantiles: 5,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 64,
        dropout_rate: 0.0,
        ..Default::default()
    };

    let mut tft = TemporalFusionTransformer::new(config.clone())?;

    // Create test inputs with distinct patterns
    let batch_size = 4;
    let historical = Tensor::randn(0.0f32, 1.0, (batch_size, 10, 64), &device)?;
    let future = Tensor::randn(0.0f32, 1.0, (batch_size, 3, 10), &device)?;

    // Test with multiple static context patterns
    let patterns = vec![
        vec![0.0f32; batch_size * 5],           // All zeros
        vec![1.0f32; batch_size * 5],           // All ones
        vec![0.5f32; batch_size * 5],           // Mid-range
        (0..batch_size * 5).map(|i| i as f32 / 10.0).collect::<Vec<_>>(), // Gradient
    ];

    let mut predictions = Vec::new();

    for (i, pattern) in patterns.iter().enumerate() {
        let static_ctx = Tensor::from_slice(pattern, (batch_size, 5), &device)?;

        let mut tft_instance = TemporalFusionTransformer::new(config.clone())?;
        let pred = tft_instance.forward(&static_ctx, &historical, &future)?;
        predictions.push(pred);

        // Verify valid output
        let pred_data = predictions[i].flatten_all()?.to_vec1::<f32>()?;
        assert!(
            pred_data.iter().all(|&x| x.is_finite()),
            "Prediction {} should be finite",
            i
        );
    }

    // Verify that different static contexts produce different predictions
    // (proves context projection layer is active)
    for i in 0..predictions.len() - 1 {
        let diff = (predictions[i + 1].clone() - predictions[i].clone())?.abs()?;
        let mean_diff = diff.mean_all()?.to_vec0::<f32>()?;

        println!("Difference between pattern {} and {}: {:.6}", i, i + 1, mean_diff);

        assert!(
            mean_diff > 0.0001,
            "Different static contexts should produce different predictions (pattern {} vs {})",
            i,
            i + 1
        );
    }

    Ok(())
}

// ============================================================================
// TEST 5: ARCHITECTURAL IMBALANCE ANALYSIS
// ============================================================================

#[test]
fn test_tft_static_vs_temporal_feature_ratio() -> Result<(), MLError> {
    let device = Device::Cpu;
    let batch_size = 2;

    // Production configuration with architectural imbalance:
    // - 5 static features
    // - 60 timesteps × 241 features = 14,460 temporal parameters
    let config = TFTConfig {
        input_dim: 241,
        hidden_dim: 64,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 5,
        sequence_length: 60,
        num_quantiles: 9,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 241,
        dropout_rate: 0.0,
        ..Default::default()
    };

    // Compute feature ratio
    let static_params = config.num_static_features;
    let temporal_params = config.sequence_length * config.num_unknown_features;
    let feature_ratio = temporal_params as f64 / static_params as f64;

    println!("Architectural analysis:");
    println!("  Static parameters:   {}", static_params);
    println!("  Temporal parameters: {}", temporal_params);
    println!("  Ratio (temporal/static): {:.1}x", feature_ratio);

    // Validate imbalance hypothesis from Wave 7.5
    assert!(
        feature_ratio > 1000.0,
        "Temporal features should dominate (ratio: {:.1}x, expected >1000x)",
        feature_ratio
    );

    // Create test inputs
    let historical = Tensor::randn(0.0f32, 1.0, (batch_size, 60, 241), &device)?;
    let future = Tensor::randn(0.0f32, 1.0, (batch_size, 5, 10), &device)?;

    // Test 1: Strong static signal vs weak temporal signal
    let static_strong = Tensor::full(5.0f32, (batch_size, 5), &device)?;
    let historical_weak = historical.clone()? * 0.1; // Scale down temporal features

    let mut tft1 = TemporalFusionTransformer::new(config.clone())?;
    let pred_static_dominant = tft1.forward(&static_strong, &historical_weak, &future)?;

    // Test 2: Weak static signal vs strong temporal signal
    let static_weak = Tensor::full(0.1f32, (batch_size, 5), &device)?;
    let historical_strong = historical.clone()? * 5.0; // Scale up temporal features

    let mut tft2 = TemporalFusionTransformer::new(config.clone())?;
    let pred_temporal_dominant = tft2.forward(&static_weak, &historical_strong, &future)?;

    // Compute prediction magnitudes
    let static_mag = pred_static_dominant.abs()?.mean_all()?.to_vec0::<f32>()?;
    let temporal_mag = pred_temporal_dominant.abs()?.mean_all()?.to_vec0::<f32>()?;

    println!("Prediction magnitudes:");
    println!("  Static-dominant:   {:.6}", static_mag);
    println!("  Temporal-dominant: {:.6}", temporal_mag);

    // Temporal features should have greater influence due to architectural imbalance
    // (This may not always hold due to Xavier initialization, but documents expected behavior)
    let magnitude_ratio = temporal_mag / static_mag.max(1e-6);
    println!("  Magnitude ratio (temporal/static): {:.2}x", magnitude_ratio);

    // Document that this test demonstrates architectural imbalance
    // but doesn't enforce strict inequality (depends on weight initialization)
    assert!(
        static_mag > 0.0 && temporal_mag > 0.0,
        "Both static and temporal features should contribute"
    );

    Ok(())
}

// ============================================================================
// TEST 6: STATIC CONTEXT SENSITIVITY ACROSS HORIZONS
// ============================================================================

#[test]
fn test_tft_static_context_horizon_sensitivity() -> Result<(), MLError> {
    let device = Device::Cpu;
    let batch_size = 2;

    let config = TFTConfig {
        input_dim: 64,
        hidden_dim: 32,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 10, // Multiple horizons to test sensitivity
        sequence_length: 20,
        num_quantiles: 5,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 64,
        dropout_rate: 0.0,
        ..Default::default()
    };

    // Create inputs
    let historical = Tensor::randn(0.0f32, 1.0, (batch_size, 20, 64), &device)?;
    let future = Tensor::randn(0.0f32, 1.0, (batch_size, 10, 10), &device)?;

    // Test with different static contexts
    let static_zero = Tensor::zeros((batch_size, 5), DType::F32, &device)?;
    let static_signal = Tensor::ones((batch_size, 5), DType::F32, &device)?;

    let mut tft_zero = TemporalFusionTransformer::new(config.clone())?;
    let pred_zero = tft_zero.forward(&static_zero, &historical, &future)?;

    let mut tft_signal = TemporalFusionTransformer::new(config.clone())?;
    let pred_signal = tft_signal.forward(&static_signal, &historical, &future)?;

    // Analyze impact across prediction horizons
    let pred_zero_data = pred_zero.to_vec3::<f32>()?;
    let pred_signal_data = pred_signal.to_vec3::<f32>()?;

    println!("Static context impact by horizon:");

    for horizon in 0..10 {
        let mut horizon_diff = 0.0f32;
        let mut count = 0;

        for batch in 0..batch_size {
            for quantile in 0..5 {
                let diff = (pred_signal_data[batch][horizon][quantile]
                           - pred_zero_data[batch][horizon][quantile]).abs();
                horizon_diff += diff;
                count += 1;
            }
        }

        let mean_horizon_diff = horizon_diff / count as f32;
        println!("  Horizon {}: mean diff = {:.6}", horizon, mean_horizon_diff);

        // Static context should affect all horizons (at least minimally)
        assert!(
            mean_horizon_diff > 0.0001 || horizon == 0,
            "Static context should affect horizon {} (diff: {:.6})",
            horizon,
            mean_horizon_diff
        );
    }

    Ok(())
}

// ============================================================================
// TEST 7: STATIC CONTEXT WITH EXTREME VALUES
// ============================================================================

#[test]
fn test_tft_static_context_extreme_values() -> Result<(), MLError> {
    let device = Device::Cpu;
    let batch_size = 2;

    let config = TFTConfig {
        input_dim: 64,
        hidden_dim: 32,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 3,
        sequence_length: 15,
        num_quantiles: 5,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 64,
        dropout_rate: 0.0,
        ..Default::default()
    };

    let historical = Tensor::randn(0.0f32, 1.0, (batch_size, 15, 64), &device)?;
    let future = Tensor::randn(0.0f32, 1.0, (batch_size, 3, 10), &device)?;

    // Test extreme static context values
    let extreme_values = vec![
        -10.0f32, // Large negative
        -1.0,     // Moderate negative
        0.0,      // Zero
        1.0,      // Moderate positive
        10.0,     // Large positive
    ];

    let mut predictions = Vec::new();

    for &value in &extreme_values {
        let static_extreme = Tensor::full(value, (batch_size, 5), &device)?;

        let mut tft = TemporalFusionTransformer::new(config.clone())?;
        let pred = tft.forward(&static_extreme, &historical, &future)?;

        // Verify predictions are finite (no explosion/vanishing)
        let pred_data = pred.flatten_all()?.to_vec1::<f32>()?;
        assert!(
            pred_data.iter().all(|&x| x.is_finite()),
            "Predictions should be finite with static context = {}",
            value
        );

        predictions.push(pred);

        let mean_pred = pred_data.iter().sum::<f32>() / pred_data.len() as f32;
        println!("Static context = {:.1}: mean prediction = {:.6}", value, mean_pred);
    }

    // Verify that extreme values produce different predictions
    for i in 0..predictions.len() - 1 {
        let diff = (predictions[i + 1].clone() - predictions[i].clone())?.abs()?;
        let mean_diff = diff.mean_all()?.to_vec0::<f32>()?;

        assert!(
            mean_diff > 0.0001,
            "Different extreme values should produce different predictions (value {} vs {})",
            extreme_values[i],
            extreme_values[i + 1]
        );
    }

    Ok(())
}