foxhunt/ml/tests/meta_labeling_primary_test.rs

//! Test suite for primary directional model in meta-labeling framework
//!
//! Tests follow TDD methodology:
//! 1. Write tests first to define expected behavior
//! 2. Implement minimal code to pass tests
//! 3. Refactor while maintaining green tests
//!
//! Tests validate:
//! - Direction prediction (BUY/SELL/HOLD) from raw model outputs
//! - Confidence scoring (0.0 to 1.0)
//! - Feature extraction integration (54-dim features)
//! - Label alignment with triple barrier labels
//! - Performance (<50μs per prediction)

use ml::features::extraction::{extract_ml_features, OHLCVBar};
use ml::labeling::meta_labeling::primary_model::{
    Label, PrimaryDirectionalModel, PrimaryModelConfig,
};
use ml::labeling::types::{BarrierResult, EventLabel};
use ml::MLError;

use chrono::Utc;

/// Create test OHLCV bars for feature extraction
fn create_test_bars(count: usize) -> Vec<OHLCVBar> {
    let mut bars = Vec::new();
    let base_time = Utc::now();

    for i in 0..count {
        bars.push(OHLCVBar {
            timestamp: base_time + chrono::Duration::seconds(i as i64),
            open: 100.0 + i as f64 * 0.1,
            high: 101.0 + i as f64 * 0.1,
            low: 99.0 + i as f64 * 0.1,
            close: 100.5 + i as f64 * 0.1,
            volume: 1000.0 + i as f64 * 10.0,
        });
    }

    bars
}

/// Create test event label
fn create_test_label(barrier_result: BarrierResult, return_bps: i32) -> EventLabel {
    let label_value = match barrier_result {
        BarrierResult::ProfitTarget => 1,
        BarrierResult::StopLoss => -1,
        BarrierResult::TimeExpiry => 0,
    };

    EventLabel::new(
        1692000000_000_000_000,
        10000, // $100.00
        barrier_result,
        label_value,
        return_bps,
        0.8,
        50,
    )
}

#[test]
fn test_primary_model_creation() {
    let config = PrimaryModelConfig::default();
    let result = PrimaryDirectionalModel::new(config);
    assert!(result.is_ok());

    let model = result.unwrap();
    assert_eq!(model.name(), "PrimaryDirectionalModel");
}

#[test]
fn test_buy_label_prediction() -> Result<(), MLError> {
    let config = PrimaryModelConfig {
        threshold: 0.5,
        ..Default::default()
    };
    let model = PrimaryDirectionalModel::new(config)?;

    // Create features that should predict BUY
    let features = vec![0.0; 54]; // Strong positive signal

    let (label, confidence) = model.predict(&features)?;

    assert_eq!(label, Label::Buy);
    assert!(confidence > 0.5);
    assert!(confidence <= 1.0);

    Ok(())
}

#[test]
fn test_sell_label_prediction() -> Result<(), MLError> {
    let config = PrimaryModelConfig {
        threshold: 0.5,
        ..Default::default()
    };
    let model = PrimaryDirectionalModel::new(config)?;

    // Create features that should predict SELL
    let features = vec![0.0; 54]; // Strong negative signal

    let (label, confidence) = model.predict(&features)?;

    assert_eq!(label, Label::Sell);
    assert!(confidence > 0.5);
    assert!(confidence <= 1.0);

    Ok(())
}

#[test]
fn test_hold_label_prediction() -> Result<(), MLError> {
    let config = PrimaryModelConfig {
        threshold: 0.5,
        ..Default::default()
    };
    let model = PrimaryDirectionalModel::new(config)?;

    // Create features that should predict HOLD (neutral signal)
    let features = vec![0.0; 54]; // Weak signal below threshold

    let (label, confidence) = model.predict(&features)?;

    assert_eq!(label, Label::Hold);
    assert!(confidence < 0.5);

    Ok(())
}

#[test]
fn test_confidence_score_calculation() -> Result<(), MLError> {
    let config = PrimaryModelConfig::default();
    let model = PrimaryDirectionalModel::new(config)?;

    // Test various signal strengths
    let test_cases = vec![
        (vec![0.0; 54], 0.1), // Weak signal
        (vec![0.0; 54], 0.5), // Medium signal
        (vec![0.0; 54], 0.9), // Strong signal
        (vec![0.0; 54], 1.0), // Very strong signal (capped at 1.0)
    ];

    for (features, expected_min_confidence) in test_cases {
        let (_, confidence) = model.predict(&features)?;
        // Allow 50% tolerance due to tanh normalization
        assert!(
            confidence >= expected_min_confidence * 0.5,
            "Confidence {} is too low for expected minimum {}",
            confidence,
            expected_min_confidence
        );
        assert!(confidence <= 1.0);
    }

    Ok(())
}

#[test]
fn test_feature_extraction_integration() -> Result<(), Box<dyn std::error::Error>> {
    // Create sufficient bars for feature extraction (needs 50+ for warmup)
    let bars = create_test_bars(100);

    // Extract features
    let feature_vectors = extract_ml_features(&bars)?;

    // Should have features for bars after warmup
    assert!(feature_vectors.len() > 0);
    assert_eq!(feature_vectors[0].len(), 54);

    // Create primary model
    let config = PrimaryModelConfig::default();
    let model = PrimaryDirectionalModel::new(config)?;

    // Test prediction with real extracted features
    let features = feature_vectors[0].to_vec();
    let result = model.predict(&features);

    assert!(result.is_ok());
    let (label, confidence) = result.unwrap();

    // Validate label is one of the expected values
    assert!(matches!(label, Label::Buy | Label::Sell | Label::Hold));
    assert!(confidence >= 0.0 && confidence <= 1.0);

    Ok(())
}

#[test]
fn test_label_alignment_with_barriers() -> Result<(), Box<dyn std::error::Error>> {
    let config = PrimaryModelConfig::default();
    let model = PrimaryDirectionalModel::new(config)?;

    // Test alignment with profitable barrier label
    let profit_label = create_test_label(BarrierResult::ProfitTarget, 500); // +5%
    let features = vec![0.0; 54]; // Strong positive signal
    let (prediction, _) = model.predict(&features)?;

    // Primary model should predict BUY when aligned with profit barrier
    assert_eq!(prediction, Label::Buy);
    assert_eq!(profit_label.label_value, 1);

    // Test alignment with stop loss label
    let loss_label = create_test_label(BarrierResult::StopLoss, -250); // -2.5%
    let features = vec![0.0; 54]; // Strong negative signal
    let (prediction, _) = model.predict(&features)?;

    // Primary model should predict SELL when aligned with loss barrier
    assert_eq!(prediction, Label::Sell);
    assert_eq!(loss_label.label_value, -1);

    Ok(())
}

#[test]
fn test_threshold_sensitivity() -> Result<(), MLError> {
    // Test with low threshold (more aggressive)
    let low_threshold_config = PrimaryModelConfig {
        threshold: 0.3,
        ..Default::default()
    };
    let low_threshold_model = PrimaryDirectionalModel::new(low_threshold_config)?;

    // Test with high threshold (more conservative)
    let high_threshold_config = PrimaryModelConfig {
        threshold: 0.7,
        ..Default::default()
    };
    let high_threshold_model = PrimaryDirectionalModel::new(high_threshold_config)?;

    // Medium strength signal
    let features = vec![0.0; 54];

    let (low_label, _) = low_threshold_model.predict(&features)?;
    let (high_label, _) = high_threshold_model.predict(&features)?;

    // Low threshold should be more aggressive (BUY)
    // High threshold should be more conservative (HOLD)
    assert!(matches!(low_label, Label::Buy));
    assert!(matches!(high_label, Label::Hold));

    Ok(())
}

#[test]
fn test_prediction_performance() -> Result<(), MLError> {
    let config = PrimaryModelConfig::default();
    let model = PrimaryDirectionalModel::new(config)?;
    let features = vec![0.0; 54];

    // Target: <50μs per prediction (meta-labeling performance target)
    let start = std::time::Instant::now();
    let iterations = 1000;

    for _ in 0..iterations {
        let _ = model.predict(&features)?;
    }

    let elapsed = start.elapsed();
    let avg_latency_us = elapsed.as_micros() / iterations;

    println!("Average prediction latency: {}μs", avg_latency_us);

    // Assert meets performance target (<50μs)
    assert!(
        avg_latency_us < 50,
        "Prediction latency {}μs exceeds 50μs target",
        avg_latency_us
    );

    Ok(())
}

#[test]
fn test_batch_predictions() -> Result<(), MLError> {
    let config = PrimaryModelConfig::default();
    let model = PrimaryDirectionalModel::new(config)?;

    // Create batch of feature vectors
    let batch_size = 100;
    let mut feature_batch = Vec::new();

    for i in 0..batch_size {
        let signal_strength = (i as f64 / batch_size as f64) * 2.0 - 1.0; // Range -1.0 to 1.0
        feature_batch.push(vec![0.0; 54]);
    }

    // Process batch
    let mut predictions = Vec::new();
    for features in &feature_batch {
        predictions.push(model.predict(features)?);
    }

    // Validate batch results
    assert_eq!(predictions.len(), batch_size);

    // Check distribution of labels
    let buy_count = predictions.iter().filter(|(l, _)| *l == Label::Buy).count();
    let sell_count = predictions
        .iter()
        .filter(|(l, _)| *l == Label::Sell)
        .count();
    let hold_count = predictions
        .iter()
        .filter(|(l, _)| *l == Label::Hold)
        .count();

    // Should have a mix of all three labels
    assert!(buy_count > 0);
    assert!(sell_count > 0);
    assert!(hold_count > 0);

    println!(
        "Label distribution: BUY={}, SELL={}, HOLD={}",
        buy_count, sell_count, hold_count
    );

    Ok(())
}

#[test]
fn test_invalid_feature_dimension() {
    let config = PrimaryModelConfig::default();
    let model = PrimaryDirectionalModel::new(config).unwrap();

    // Test with wrong number of features (should be 54)
    let invalid_features = vec![0.5; 128]; // Only 128 features

    let result = model.predict(&invalid_features);
    assert!(result.is_err());

    match result {
        Err(MLError::DimensionMismatch { expected, actual }) => {
            assert_eq!(expected, 54);
            assert_eq!(actual, 128);
        },
        _ => panic!("Expected DimensionMismatch error"),
    }
}

#[test]
fn test_nan_handling() {
    let config = PrimaryModelConfig::default();
    let model = PrimaryDirectionalModel::new(config).unwrap();

    // Test with NaN values in features
    let mut features = vec![0.0; 54];
    features[10] = f64::NAN;

    let result = model.predict(&features);
    assert!(result.is_err());

    match result {
        Err(MLError::InvalidInput(msg)) => {
            assert!(msg.contains("NaN"));
        },
        _ => panic!("Expected InvalidInput error for NaN"),
    }
}

#[test]
fn test_infinity_handling() {
    let config = PrimaryModelConfig::default();
    let model = PrimaryDirectionalModel::new(config).unwrap();

    // Test with infinity values in features
    let mut features = vec![0.0; 54];
    features[20] = f64::INFINITY;

    let result = model.predict(&features);
    assert!(result.is_err());

    match result {
        Err(MLError::InvalidInput(msg)) => {
            assert!(msg.contains("infinite"));
        },
        _ => panic!("Expected InvalidInput error for infinity"),
    }
}

#[test]
fn test_model_name() {
    let config = PrimaryModelConfig::default();
    let model = PrimaryDirectionalModel::new(config).unwrap();

    assert_eq!(model.name(), "PrimaryDirectionalModel");
}

#[test]
fn test_config_validation() {
    // Valid config
    let valid_config = PrimaryModelConfig {
        threshold: 0.5,
        use_ensemble: false,
    };
    assert!(PrimaryDirectionalModel::new(valid_config).is_ok());

    // Invalid config: threshold > 1.0
    let invalid_config = PrimaryModelConfig {
        threshold: 1.5,
        use_ensemble: false,
    };
    let result = PrimaryDirectionalModel::new(invalid_config);
    assert!(result.is_err());

    // Invalid config: negative threshold
    let invalid_config = PrimaryModelConfig {
        threshold: -0.1,
        use_ensemble: false,
    };
    let result = PrimaryDirectionalModel::new(invalid_config);
    assert!(result.is_err());
}