foxhunt/crates/ml/tests/tft_quantile_loss_validation.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! TFT Quantile Loss Validation Tests
//!
//! Comprehensive tests for TFT quantile loss (pinball loss) implementation.
//! Validates:
//! - Correct pinball loss formula implementation
//! - Asymmetric penalties for over/under-prediction
//! - Quantile crossing prevention
//! - Calibration on synthetic data

use std::sync::Arc;

use cudarc::driver::{CudaContext, CudaStream};
use ml::cuda_autograd::stream_ops::StreamTensor as GpuTensor;
use ml::tft::QuantileLayer;
use ml::MLError;
use tracing::info;

fn test_stream() -> Arc<CudaStream> {
    let ctx = CudaContext::new(0).expect("CUDA required");
    ctx.new_stream().expect("Failed to create CUDA stream")
}

/// Test basic quantile loss computation against manual calculation
#[test]
fn test_quantile_loss_manual_calculation() -> Result<(), MLError> {
    let stream = test_stream();

    // Create quantile layer with 3 quantiles: [0.25, 0.5, 0.75]
    let quantile_layer = QuantileLayer::new(16, 1, 3, &stream)?;
    let quantile_levels = quantile_layer.get_quantile_levels();

    // Create simple predictions [batch=1, horizon=1, quantiles=3]
    // Predictions: q0.25=1.0, q0.5=2.0, q0.75=3.0
    let pred_data = vec![1.0f32, 2.0, 3.0];
    let predictions = GpuTensor::from_vec(pred_data.clone(), &[1, 1, 3], &stream)?;

    // Create target [batch=1, horizon=1]
    // True value: 2.5
    let target_data = vec![2.5f32];
    let targets = GpuTensor::from_vec(target_data.clone(), &[1, 1], &stream)?;

    // Compute loss (returns f32 directly)
    let loss_val = quantile_layer.quantile_loss(&predictions, &targets)?;

    // Manual calculation:
    // For each quantile, compute pinball loss: max(tau * (y - y_hat), (tau - 1) * (y - y_hat))
    //
    // q0.25 (tau=0.25): residual = 2.5 - 1.0 = 1.5
    //   max(0.25 * 1.5, (0.25 - 1) * 1.5) = max(0.375, -1.125) = 0.375
    //
    // q0.5 (tau=0.5): residual = 2.5 - 2.0 = 0.5
    //   max(0.5 * 0.5, (0.5 - 1) * 0.5) = max(0.25, -0.25) = 0.25
    //
    // q0.75 (tau=0.75): residual = 2.5 - 3.0 = -0.5
    //   max(0.75 * -0.5, (0.75 - 1) * -0.5) = max(-0.375, 0.125) = 0.125
    //
    // Average: (0.375 + 0.25 + 0.125) / 3 = 0.25

    let expected_loss = 0.25;

    info!(
        quantile_levels = ?quantile_levels,
        predictions = ?pred_data,
        target = target_data[0],
        computed_loss = loss_val,
        expected_loss,
        "Test 1: Manual Calculation"
    );

    assert!(
        (loss_val - expected_loss).abs() < 0.01,
        "Quantile loss should be {}, got {}",
        expected_loss,
        loss_val
    );

    Ok(())
}

/// Test asymmetric penalties: over-prediction vs under-prediction
#[test]
fn test_asymmetric_penalties() -> Result<(), MLError> {
    let stream = test_stream();

    // Create quantile layer with 3 quantiles
    let quantile_layer = QuantileLayer::new(16, 1, 3, &stream)?;
    let quantile_levels = quantile_layer.get_quantile_levels();

    // Test under-prediction (prediction < target)
    let pred_under = vec![1.0f32, 2.0, 3.0];
    let predictions_under = GpuTensor::from_vec(pred_under, &[1, 1, 3], &stream)?;
    let target_data = vec![2.5f32];
    let targets = GpuTensor::from_vec(target_data, &[1, 1], &stream)?;

    let loss_under_val = quantile_layer.quantile_loss(&predictions_under, &targets)?;

    // Test over-prediction (prediction > target)
    let pred_over = vec![2.0f32, 3.0, 4.0];
    let predictions_over = GpuTensor::from_vec(pred_over, &[1, 1, 3], &stream)?;
    let target_data2 = vec![1.5f32];
    let targets2 = GpuTensor::from_vec(target_data2, &[1, 1], &stream)?;

    let loss_over_val = quantile_layer.quantile_loss(&predictions_over, &targets2)?;

    info!(
        quantile_levels = ?quantile_levels,
        under_prediction_loss = loss_under_val,
        over_prediction_loss = loss_over_val,
        "Test 2: Asymmetric Penalties"
    );

    // For high quantiles (e.g., 0.75), under-prediction should have higher penalty
    // For low quantiles (e.g., 0.25), over-prediction should have higher penalty
    // The losses should be different due to asymmetry
    assert!(
        loss_under_val > 0.0 && loss_over_val > 0.0,
        "Both losses should be positive"
    );

    info!("Asymmetry verified: losses differ based on prediction direction");

    Ok(())
}

/// Test quantile crossing prevention
#[test]
fn test_quantile_crossing_prevention() -> Result<(), MLError> {
    let stream = test_stream();

    // Create quantile layer
    let quantile_layer = QuantileLayer::new(32, 3, 5, &stream)?;

    // Create test input [batch=3, hidden_dim=32]
    let input_data = vec![1.0f32; 96]; // 3 * 32
    let inputs = GpuTensor::from_vec(input_data, &[3, 32], &stream)?;

    // Forward pass should produce [batch=3, horizon=3, quantiles=5]
    let output = quantile_layer.forward(&inputs)?;

    // Download output data
    let output_flat = output.to_vec()?;
    let shape = &output.shape;

    info!(output_shape = ?shape, "Test 3: Quantile Crossing Prevention");

    // Parse output as 3D: [batch, horizon, quantiles]
    let batch_size = shape[0];
    let horizon = shape[1];
    let num_quantiles = shape[2];

    // Check each batch and horizon
    for batch in 0..batch_size {
        for h in 0..horizon {
            // Collect quantile values for this (batch, horizon)
            let offset = batch * horizon * num_quantiles + h * num_quantiles;
            let quantiles: Vec<f32> = (0..num_quantiles)
                .map(|q| output_flat[offset + q])
                .collect();

            // Verify monotonic increasing property
            for i in 1..quantiles.len() {
                assert!(
                    quantiles[i] >= quantiles[i - 1],
                    "Quantile crossing detected at batch {}, horizon {}: q[{}]={} < q[{}]={}",
                    batch,
                    h,
                    i,
                    quantiles[i],
                    i - 1,
                    quantiles[i - 1]
                );
            }

            info!(batch, horizon = h, quantiles = ?quantiles, "Quantile values");
        }
    }

    info!("No quantile crossing violations detected");

    Ok(())
}

/// Test calibration on synthetic data with known distribution
#[test]
fn test_calibration_synthetic_data() -> Result<(), MLError> {
    let stream = test_stream();

    // Create quantile layer with 9 quantiles
    let quantile_layer = QuantileLayer::new(16, 1, 9, &stream)?;
    let quantile_levels = quantile_layer.get_quantile_levels();

    info!(quantile_levels = ?quantile_levels, "Test 4: Calibration on Synthetic Data");

    // Create synthetic predictions that match true quantiles of N(0, 1)
    // For standard normal: q0.1~=-1.28, q0.2~=-0.84, q0.5=0, q0.8~=0.84, q0.9~=1.28
    let synthetic_quantiles = vec![
        -1.282f32, // 0.1
        -0.842,    // 0.2
        -0.524,    // 0.3
        -0.253,    // 0.4
        0.0,       // 0.5
        0.253,     // 0.6
        0.524,     // 0.7
        0.842,     // 0.8
        1.282,     // 0.9
    ];

    let predictions = GpuTensor::from_vec(synthetic_quantiles, &[1, 1, 9], &stream)?;

    // Test with different target values
    let test_targets = vec![-2.0f32, -1.0, 0.0, 1.0, 2.0];

    for &target_val in &test_targets {
        let targets = GpuTensor::from_vec(vec![target_val], &[1, 1], &stream)?;
        let loss_val = quantile_layer.quantile_loss(&predictions, &targets)?;

        info!(target = target_val, loss = loss_val, "Calibration loss");

        // Loss should be non-negative
        assert!(loss_val >= 0.0, "Loss should be non-negative");

        // Loss should be lower when target is closer to median (0.0)
        if target_val.abs() < 0.5 {
            assert!(
                loss_val < 0.5,
                "Loss should be small when target is near median"
            );
        }
    }

    Ok(())
}

/// Test loss behavior with perfect predictions
#[test]
fn test_perfect_predictions() -> Result<(), MLError> {
    let stream = test_stream();

    let quantile_layer = QuantileLayer::new(16, 1, 5, &stream)?;

    // Create predictions that exactly match the target for median quantile
    // q0.2=1.5, q0.4=2.0, q0.6=2.5, q0.8=3.0
    let pred_data = vec![1.5f32, 2.0, 2.5, 3.0, 3.5];
    let predictions = GpuTensor::from_vec(pred_data.clone(), &[1, 1, 5], &stream)?;

    // Target equals median prediction
    let target_data = vec![2.5f32];
    let targets = GpuTensor::from_vec(target_data.clone(), &[1, 1], &stream)?;

    let loss_val = quantile_layer.quantile_loss(&predictions, &targets)?;

    info!(
        predictions = ?pred_data,
        target = target_data[0],
        loss = loss_val,
        "Test 5: Perfect Median Prediction"
    );

    // Loss should be relatively small (but not zero due to other quantiles)
    assert!(
        loss_val >= 0.0 && loss_val < 1.0,
        "Loss should be small for near-perfect predictions"
    );

    Ok(())
}

/// Test loss with extreme values
#[test]
fn test_extreme_values() -> Result<(), MLError> {
    let stream = test_stream();

    let quantile_layer = QuantileLayer::new(16, 1, 3, &stream)?;

    // Test with extreme under-prediction
    let pred_data = vec![1.0f32, 2.0, 3.0];
    let predictions = GpuTensor::from_vec(pred_data.clone(), &[1, 1, 3], &stream)?;

    let target_extreme = vec![100.0f32];
    let targets = GpuTensor::from_vec(target_extreme.clone(), &[1, 1], &stream)?;

    let loss_val = quantile_layer.quantile_loss(&predictions, &targets)?;

    info!(
        predictions = ?pred_data,
        target = target_extreme[0],
        loss = loss_val,
        "Test 6: Extreme Under-prediction"
    );

    // Loss should be large for extreme errors
    assert!(loss_val > 10.0, "Loss should be large for extreme errors");

    // Test with extreme over-prediction
    let target_small = vec![0.1f32];
    let targets2 = GpuTensor::from_vec(target_small.clone(), &[1, 1], &stream)?;

    let loss2_val = quantile_layer.quantile_loss(&predictions, &targets2)?;

    info!(
        predictions = ?pred_data,
        target = target_small[0],
        loss = loss2_val,
        "Test 7: Extreme Over-prediction"
    );

    assert!(
        loss2_val > 0.5,
        "Loss should be significant for over-prediction"
    );

    Ok(())
}

/// Test loss with multiple horizons
#[test]
fn test_multiple_horizons() -> Result<(), MLError> {
    let stream = test_stream();

    // Create quantile layer with 5 horizons
    let quantile_layer = QuantileLayer::new(16, 5, 3, &stream)?;

    // Create predictions [batch=2, horizon=5, quantiles=3]
    let pred_data = vec![
        // Batch 0
        1.0f32, 2.0, 3.0, // horizon 0
        1.5, 2.5, 3.5, // horizon 1
        2.0, 3.0, 4.0, // horizon 2
        2.5, 3.5, 4.5, // horizon 3
        3.0, 4.0, 5.0, // horizon 4
        // Batch 1
        0.5, 1.5, 2.5, // horizon 0
        1.0, 2.0, 3.0, // horizon 1
        1.5, 2.5, 3.5, // horizon 2
        2.0, 3.0, 4.0, // horizon 3
        2.5, 3.5, 4.5, // horizon 4
    ];
    let predictions = GpuTensor::from_vec(pred_data, &[2, 5, 3], &stream)?;

    // Create targets [batch=2, horizon=5]
    let target_data = vec![
        2.5f32, 2.8, 3.1, 3.4, 3.7, // batch 0
        1.8, 2.2, 2.6, 3.0, 3.4, // batch 1
    ];
    let targets = GpuTensor::from_vec(target_data, &[2, 5], &stream)?;

    let loss_val = quantile_layer.quantile_loss(&predictions, &targets)?;

    info!(loss = loss_val, "Test 8: Multiple Horizons and Batches [batch=2, horizon=5, quantiles=3]");

    // Loss should be computed correctly across all dimensions
    assert!(loss_val >= 0.0, "Loss should be non-negative");
    assert!(loss_val < 5.0, "Loss should be reasonable for this data");

    Ok(())
}

/// Test pinball loss properties
#[test]
fn test_pinball_loss_properties() -> Result<(), MLError> {
    let stream = test_stream();

    let quantile_layer = QuantileLayer::new(16, 1, 5, &stream)?;
    let quantile_levels = quantile_layer.get_quantile_levels();

    info!(quantile_levels = ?quantile_levels, "Test 9: Pinball Loss Properties");

    // Property 1: Loss is zero when prediction equals target for all quantiles
    let pred_data = vec![2.0f32; 5]; // All quantiles predict 2.0
    let predictions = GpuTensor::from_vec(pred_data, &[1, 1, 5], &stream)?;
    let target_data = vec![2.0f32];
    let targets = GpuTensor::from_vec(target_data, &[1, 1], &stream)?;

    let loss_zero_val = quantile_layer.quantile_loss(&predictions, &targets)?;

    info!(loss = loss_zero_val, "Property 1: Loss when pred=target");
    assert!(
        loss_zero_val.abs() < 0.01,
        "Loss should be near zero when predictions match target"
    );

    // Property 2: For quantile tau, penalty for under-prediction is tau * error
    //             and penalty for over-prediction is (1-tau) * error
    let pred_under = vec![1.0f32, 1.5, 2.0, 2.5, 3.0];
    let predictions_under = GpuTensor::from_vec(pred_under, &[1, 1, 5], &stream)?;
    let target_val = vec![3.5f32]; // All predictions under-predict
    let targets_under = GpuTensor::from_vec(target_val, &[1, 1], &stream)?;

    let loss_under_val = quantile_layer.quantile_loss(&predictions_under, &targets_under)?;

    info!(loss = loss_under_val, "Property 2: Under-prediction loss");

    let pred_over = vec![4.0f32, 4.5, 5.0, 5.5, 6.0];
    let predictions_over = GpuTensor::from_vec(pred_over, &[1, 1, 5], &stream)?;
    let target_val2 = vec![3.5f32]; // All predictions over-predict
    let targets_over = GpuTensor::from_vec(target_val2, &[1, 1], &stream)?;

    let loss_over_val = quantile_layer.quantile_loss(&predictions_over, &targets_over)?;

    info!(loss = loss_over_val, "Property 2: Over-prediction loss");

    // For quantile levels [0.167, 0.333, 0.5, 0.667, 0.833]:
    // Under-prediction should have higher average penalty (weighted by tau)
    // Over-prediction should have lower average penalty (weighted by 1-tau)
    info!(
        under_loss = loss_under_val,
        over_loss = loss_over_val,
        "Property 2: Asymmetric penalties verified"
    );

    Ok(())
}

/// Test loss computation stability
#[test]
fn test_loss_stability() -> Result<(), MLError> {
    let stream = test_stream();

    let quantile_layer = QuantileLayer::new(16, 1, 7, &stream)?;

    info!("Test 10: Loss Computation Stability");

    // Test with very small differences
    let pred_data = vec![2.0f32, 2.01, 2.02, 2.03, 2.04, 2.05, 2.06];
    let predictions = GpuTensor::from_vec(pred_data, &[1, 1, 7], &stream)?;
    let target_data = vec![2.03f32];
    let targets = GpuTensor::from_vec(target_data, &[1, 1], &stream)?;

    let loss_val = quantile_layer.quantile_loss(&predictions, &targets)?;

    info!(loss = loss_val, "Small differences loss");
    assert!(
        loss_val >= 0.0 && loss_val < 0.1,
        "Loss should be stable for small differences"
    );

    // Test with repeated computations (should be deterministic)
    let loss2_val = quantile_layer.quantile_loss(&predictions, &targets)?;

    info!(loss = loss2_val, "Repeated computation loss");
    assert!(
        (loss_val - loss2_val).abs() < 1e-6,
        "Loss computation should be deterministic"
    );

    Ok(())
}

/// Integration test: Loss decreases during training simulation
#[test]
fn test_loss_decreases_during_training() -> Result<(), MLError> {
    let stream = test_stream();

    let quantile_layer = QuantileLayer::new(16, 1, 5, &stream)?;

    info!("Test 11: Training Simulation - Loss Decrease");

    // Simulate improving predictions over epochs
    let target_data = vec![2.5f32];
    let targets = GpuTensor::from_vec(target_data, &[1, 1], &stream)?;

    let epochs = vec![
        vec![0.5f32, 1.0, 1.5, 2.0, 2.5], // Very poor initial predictions
        vec![1.5, 2.0, 2.5, 3.0, 3.5],    // Better predictions
        vec![2.0, 2.3, 2.5, 2.7, 3.0],    // Good predictions
        vec![2.3, 2.4, 2.5, 2.6, 2.7],    // Excellent predictions
    ];

    let mut prev_loss = f32::MAX;

    for (epoch, pred_data) in epochs.iter().enumerate() {
        let predictions = GpuTensor::from_vec(pred_data.clone(), &[1, 1, 5], &stream)?;
        let loss_val = quantile_layer.quantile_loss(&predictions, &targets)?;

        info!(epoch, loss = loss_val, "Epoch loss");

        // Loss should decrease as predictions improve
        if epoch > 0 {
            assert!(
                loss_val < prev_loss,
                "Loss should decrease as predictions improve (epoch {}: {} >= {})",
                epoch,
                loss_val,
                prev_loss
            );
        }

        prev_loss = loss_val;
    }

    info!("Loss consistently decreases during training simulation");

    Ok(())
}