foxhunt/ml/tests/sample_weights_test.rs

//! Sample Weights Test Suite (TDD)
//!
//! Tests for sample weight calculation to address:
//! - Label imbalance (buy/sell/hold distribution)
//! - Temporal decay (recent samples weighted higher)
//! - Numerical stability (normalized weights)
//!
//! Based on MLFinLab methodology for reducing overfitting

use chrono::{DateTime, Duration, Utc};

// We'll import from the module we're about to create
use ml::features::sample_weights::{SampleWeightCalculator, WeightingScheme};
use ml::labeling::meta_labeling::primary_model::Label;

/// Test helper: create timestamps with specified day offsets from now
fn create_timestamps(day_offsets: Vec<i64>) -> Vec<DateTime<Utc>> {
    let base_time = Utc::now();
    day_offsets
        .into_iter()
        .map(|offset| base_time - Duration::days(offset))
        .collect()
}

#[test]
fn test_temporal_decay_only() {
    // Test temporal decay without label balancing
    let calculator = SampleWeightCalculator::new(
        0.95,                           // decay_factor
        WeightingScheme::TemporalDecay, // scheme
    );

    // Create labels (all Buy, so no label imbalance effect)
    let labels = vec![Label::Buy; 5];

    // Create timestamps: 4 days ago, 3 days ago, ..., today
    let timestamps = create_timestamps(vec![4, 3, 2, 1, 0]);

    let weights = calculator
        .calculate(&labels, &timestamps)
        .expect("Weight calculation should succeed");

    // Verify weights are normalized (sum to 1.0)
    let sum: f64 = weights.iter().sum();
    assert!(
        (sum - 1.0).abs() < 1e-6,
        "Weights should sum to 1.0, got {}",
        sum
    );

    // Verify temporal decay pattern: more recent samples have higher weights
    assert!(
        weights[0] < weights[4],
        "Oldest sample ({}) should have lower weight than newest ({})",
        weights[0],
        weights[4]
    );

    // Verify exponential decay relationship
    // decay_factor^1 = 0.95, so weight ratios should approximately match
    for i in 0..weights.len() - 1 {
        let ratio = weights[i + 1] / weights[i];
        assert!(
            (ratio - 1.0 / 0.95).abs() < 0.01,
            "Adjacent weight ratio should be ~1.053, got {}",
            ratio
        );
    }
}

#[test]
fn test_label_balancing_only() {
    // Test label balancing without temporal decay
    let calculator = SampleWeightCalculator::new(
        1.0,                             // No decay (decay_factor = 1.0)
        WeightingScheme::LabelBalancing, // scheme
    );

    // Create imbalanced labels: 3 Buy, 1 Sell, 1 Hold
    let labels = vec![Label::Buy, Label::Buy, Label::Buy, Label::Sell, Label::Hold];

    // All timestamps the same (no temporal effect)
    let timestamps = vec![Utc::now(); 5];

    let weights = calculator
        .calculate(&labels, &timestamps)
        .expect("Weight calculation should succeed");

    // Verify weights are normalized
    let sum: f64 = weights.iter().sum();
    assert!(
        (sum - 1.0).abs() < 1e-6,
        "Weights should sum to 1.0, got {}",
        sum
    );

    // Buy appears 3 times, so each Buy sample gets 1/3 weight factor
    // Sell appears 1 time, so Sell sample gets 1/1 = 1 weight factor
    // Hold appears 1 time, so Hold sample gets 1/1 = 1 weight factor
    // After normalization, Sell and Hold should have higher weights than Buy

    let buy_weight = weights[0]; // First Buy sample
    let sell_weight = weights[3]; // Sell sample
    let hold_weight = weights[4]; // Hold sample

    assert!(
        sell_weight > buy_weight,
        "Sell (rare) should have higher weight than Buy (common)"
    );
    assert!(
        hold_weight > buy_weight,
        "Hold (rare) should have higher weight than Buy (common)"
    );

    // Sell and Hold should have approximately equal weights (both appear once)
    assert!(
        (sell_weight - hold_weight).abs() < 1e-6,
        "Sell and Hold should have equal weights (both appear once)"
    );
}

#[test]
fn test_combined_weighting() {
    // Test combining temporal decay and label balancing
    let calculator = SampleWeightCalculator::new(
        0.95,                      // decay_factor
        WeightingScheme::Combined, // Both temporal and label balancing
    );

    // Create imbalanced labels with temporal spread
    let labels = vec![
        Label::Buy,  // 4 days ago
        Label::Buy,  // 3 days ago
        Label::Sell, // 2 days ago
        Label::Hold, // 1 day ago
        Label::Buy,  // today
    ];

    let timestamps = create_timestamps(vec![4, 3, 2, 1, 0]);

    let weights = calculator
        .calculate(&labels, &timestamps)
        .expect("Weight calculation should succeed");

    // Verify normalization
    let sum: f64 = weights.iter().sum();
    assert!(
        (sum - 1.0).abs() < 1e-6,
        "Weights should sum to 1.0, got {}",
        sum
    );

    // Buy appears 3 times (indices 0, 1, 4)
    // Sell appears 1 time (index 2)
    // Hold appears 1 time (index 3)

    // The most recent Buy (index 4) should have higher weight than oldest Buy (index 0)
    assert!(
        weights[4] > weights[0],
        "Most recent Buy should have higher weight than oldest Buy"
    );

    // Recent Sell (index 2) should have high weight (recent + rare)
    // Recent Hold (index 3) should have high weight (recent + rare)
    // These two should be among the highest weights
    assert!(
        weights[2] > weights[0],
        "Recent Sell should have higher weight than old Buy"
    );
    assert!(
        weights[3] > weights[0],
        "Recent Hold should have higher weight than old Buy"
    );
}

#[test]
fn test_numerical_stability_large_time_gaps() {
    // Test with large time gaps to ensure numerical stability
    let calculator = SampleWeightCalculator::new(0.95, WeightingScheme::TemporalDecay);

    let labels = vec![Label::Buy; 3];
    // Very old sample (365 days ago), medium (30 days), recent (1 day)
    let timestamps = create_timestamps(vec![365, 30, 1]);

    let weights = calculator
        .calculate(&labels, &timestamps)
        .expect("Weight calculation should succeed");

    // Verify normalization
    let sum: f64 = weights.iter().sum();
    assert!(
        (sum - 1.0).abs() < 1e-6,
        "Weights should sum to 1.0 even with large time gaps, got {}",
        sum
    );

    // Verify all weights are positive
    for (i, &weight) in weights.iter().enumerate() {
        assert!(
            weight > 0.0,
            "Weight at index {} should be positive, got {}",
            i,
            weight
        );
    }

    // Very old sample should have negligible weight compared to recent
    assert!(
        weights[0] < weights[2] * 0.001,
        "Very old sample should have negligible weight compared to recent"
    );
}

#[test]
fn test_numerical_stability_equal_labels() {
    // Test with perfectly balanced labels
    let calculator = SampleWeightCalculator::new(1.0, WeightingScheme::LabelBalancing);

    // Equal distribution: 3 Buy, 3 Sell, 3 Hold
    let labels = vec![
        Label::Buy,
        Label::Sell,
        Label::Hold,
        Label::Buy,
        Label::Sell,
        Label::Hold,
        Label::Buy,
        Label::Sell,
        Label::Hold,
    ];

    let timestamps = vec![Utc::now(); 9];

    let weights = calculator
        .calculate(&labels, &timestamps)
        .expect("Weight calculation should succeed");

    // With equal labels and no temporal decay, all weights should be equal
    let expected_weight = 1.0 / 9.0;
    for (i, &weight) in weights.iter().enumerate() {
        assert!(
            (weight - expected_weight).abs() < 1e-6,
            "Weight at index {} should be {}, got {}",
            i,
            expected_weight,
            weight
        );
    }
}

#[test]
fn test_numerical_stability_single_sample() {
    // Edge case: single sample
    let calculator = SampleWeightCalculator::new(0.95, WeightingScheme::Combined);

    let labels = vec![Label::Buy];
    let timestamps = vec![Utc::now()];

    let weights = calculator
        .calculate(&labels, &timestamps)
        .expect("Weight calculation should succeed");

    // Single sample should have weight 1.0
    assert_eq!(weights.len(), 1);
    assert!(
        (weights[0] - 1.0).abs() < 1e-6,
        "Single sample should have weight 1.0, got {}",
        weights[0]
    );
}

#[test]
fn test_empty_input_error() {
    // Test error handling for empty inputs
    let calculator = SampleWeightCalculator::new(0.95, WeightingScheme::Combined);

    let labels = vec![];
    let timestamps = vec![];

    let result = calculator.calculate(&labels, &timestamps);

    assert!(result.is_err(), "Empty input should return an error");
}

#[test]
fn test_mismatched_lengths_error() {
    // Test error handling for mismatched input lengths
    let calculator = SampleWeightCalculator::new(0.95, WeightingScheme::Combined);

    let labels = vec![Label::Buy, Label::Sell];
    let timestamps = vec![Utc::now()]; // Only 1 timestamp for 2 labels

    let result = calculator.calculate(&labels, &timestamps);

    assert!(
        result.is_err(),
        "Mismatched input lengths should return an error"
    );
}

#[test]
fn test_invalid_decay_factor_error() {
    // Test that decay factor must be positive
    // This should panic or return error during construction

    // Test decay_factor = 0 (invalid)
    let calculator = SampleWeightCalculator::new(0.0, WeightingScheme::TemporalDecay);

    let labels = vec![Label::Buy];
    let timestamps = vec![Utc::now()];

    let result = calculator.calculate(&labels, &timestamps);
    assert!(result.is_err(), "Decay factor 0.0 should produce an error");

    // Test decay_factor > 1.0 (unusual but mathematically valid - future weighted higher)
    let calculator = SampleWeightCalculator::new(1.5, WeightingScheme::TemporalDecay);

    let result = calculator.calculate(&labels, &timestamps);
    // Should succeed (mathematically valid, just unusual)
    assert!(
        result.is_ok(),
        "Decay factor > 1.0 should be allowed (future-weighted)"
    );
}

#[test]
fn test_weights_non_negative() {
    // Ensure all weights are non-negative in all schemes
    let schemes = vec![
        WeightingScheme::TemporalDecay,
        WeightingScheme::LabelBalancing,
        WeightingScheme::Combined,
    ];

    let labels = vec![Label::Buy, Label::Sell, Label::Hold, Label::Buy];
    let timestamps = create_timestamps(vec![3, 2, 1, 0]);

    for scheme in schemes {
        let calculator = SampleWeightCalculator::new(0.95, scheme);
        let weights = calculator
            .calculate(&labels, &timestamps)
            .expect("Weight calculation should succeed");

        for (i, &weight) in weights.iter().enumerate() {
            assert!(
                weight >= 0.0,
                "Weight at index {} should be non-negative, got {}",
                i,
                weight
            );
        }
    }
}

#[test]
fn test_extreme_imbalance() {
    // Test with extreme label imbalance (99:1 ratio)
    let calculator = SampleWeightCalculator::new(1.0, WeightingScheme::LabelBalancing);

    // 99 Buy labels, 1 Sell label
    let mut labels = vec![Label::Buy; 99];
    labels.push(Label::Sell);

    let timestamps = vec![Utc::now(); 100];

    let weights = calculator
        .calculate(&labels, &timestamps)
        .expect("Weight calculation should succeed");

    // Verify normalization
    let sum: f64 = weights.iter().sum();
    assert!(
        (sum - 1.0).abs() < 1e-6,
        "Weights should sum to 1.0, got {}",
        sum
    );

    // The single Sell should have much higher weight than any Buy
    let sell_weight = weights[99];
    let buy_weight = weights[0];

    assert!(
        sell_weight > buy_weight * 50.0,
        "Rare Sell should have 50x+ weight compared to common Buy"
    );

    // Total weight for all Sell samples should roughly equal total weight for all Buy samples
    let total_sell_weight = sell_weight;
    let total_buy_weight: f64 = weights[0..99].iter().sum();

    assert!(
        (total_sell_weight - total_buy_weight).abs() < 0.1,
        "Total weight for Sell should approximately equal total weight for Buy (balanced classes)"
    );
}