foxhunt/ml/tests/imbalance_bars_test.rs

/// TDD Tests for Imbalance Bars Implementation
///
/// Imbalance bars emit when cumulative buy/sell imbalance exceeds threshold.
/// Expected improvement: +15-20% Sharpe vs time bars due to better information capture.

#[cfg(test)]
mod imbalance_bars_tests {
    use chrono::{DateTime, Utc};

    // Implemented in ml/src/features/alternative_bars.rs
    use ml::features::alternative_bars::ImbalanceBarSampler;

    fn timestamp(secs: i64) -> DateTime<Utc> {
        DateTime::from_timestamp(secs, 0).unwrap()
    }

    #[test]
    fn test_buy_tick_classification() {
        // Buy tick: price increases
        let mut sampler = ImbalanceBarSampler::new(100.0, 100.0, timestamp(0));

        // First tick at $100 (no direction yet)
        assert!(sampler.update(100.0, 10.0, timestamp(0)).is_none());

        // Price increases to $101 -> buy tick
        assert!(sampler.update(101.0, 10.0, timestamp(1)).is_none());

        // Verify imbalance increased (buy side)
        assert!(
            sampler.get_imbalance() > 0.0,
            "Buy tick should increase imbalance"
        );
    }

    #[test]
    fn test_sell_tick_classification() {
        // Sell tick: price decreases
        let mut sampler = ImbalanceBarSampler::new(100.0, 100.0, timestamp(0));

        // First tick at $100
        assert!(sampler.update(100.0, 10.0, timestamp(0)).is_none());

        // Price decreases to $99 -> sell tick
        assert!(sampler.update(99.0, 10.0, timestamp(1)).is_none());

        // Verify imbalance decreased (sell side)
        assert!(
            sampler.get_imbalance() < 0.0,
            "Sell tick should decrease imbalance"
        );
    }

    #[test]
    fn test_cumulative_imbalance_calculation() {
        let mut sampler = ImbalanceBarSampler::new(100.0, 1000.0, timestamp(0));

        // Sequence of ticks with known direction
        sampler.update(100.0, 10.0, timestamp(0)); // Baseline
        sampler.update(101.0, 20.0, timestamp(1)); // Buy: +20
        sampler.update(102.0, 15.0, timestamp(2)); // Buy: +15
        sampler.update(101.0, 10.0, timestamp(3)); // Sell: -10
        sampler.update(102.0, 25.0, timestamp(4)); // Buy: +25

        // Expected cumulative imbalance: +20 +15 -10 +25 = +50
        let imbalance = sampler.get_imbalance();
        assert!(
            (imbalance - 50.0).abs() < 0.01,
            "Cumulative imbalance should be +50, got {}",
            imbalance
        );
    }

    #[test]
    fn test_bar_formation_at_positive_threshold() {
        // Threshold = 100, emit bar when imbalance >= 100
        let mut sampler = ImbalanceBarSampler::new(100.0, 100.0, timestamp(0));

        sampler.update(100.0, 10.0, timestamp(0)); // Baseline
        assert!(sampler.update(101.0, 50.0, timestamp(1)).is_none()); // +50
        assert!(sampler.update(102.0, 40.0, timestamp(2)).is_none()); // +90

        // Next buy tick should trigger bar (90 + 20 = 110 >= 100)
        let bar = sampler.update(103.0, 20.0, timestamp(3));
        assert!(
            bar.is_some(),
            "Bar should emit when imbalance exceeds threshold"
        );

        let bar = bar.unwrap();
        assert_eq!(bar.open, 100.0);
        assert_eq!(bar.high, 103.0);
        assert_eq!(bar.low, 100.0);
        assert_eq!(bar.close, 103.0);
        assert_eq!(bar.volume, 120.0); // 10+50+40+20

        // Imbalance should reset after bar emission
        assert_eq!(sampler.get_imbalance(), 0.0);
    }

    #[test]
    fn test_bar_formation_at_negative_threshold() {
        // Test sell-side imbalance triggering bar
        let mut sampler = ImbalanceBarSampler::new(100.0, 100.0, timestamp(0));

        sampler.update(100.0, 10.0, timestamp(0)); // Baseline
        assert!(sampler.update(99.0, 50.0, timestamp(1)).is_none()); // -50
        assert!(sampler.update(98.0, 40.0, timestamp(2)).is_none()); // -90

        // Next sell tick should trigger bar (-90 - 20 = -110, abs >= 100)
        let bar = sampler.update(97.0, 20.0, timestamp(3));
        assert!(
            bar.is_some(),
            "Bar should emit when negative imbalance exceeds threshold"
        );

        let bar = bar.unwrap();
        assert_eq!(bar.open, 100.0);
        assert_eq!(bar.high, 100.0);
        assert_eq!(bar.low, 97.0);
        assert_eq!(bar.close, 97.0);
    }

    #[test]
    fn test_balanced_market_no_bar() {
        // Balanced buy/sell should not trigger bars
        let mut sampler = ImbalanceBarSampler::new(100.0, 100.0, timestamp(0));

        sampler.update(100.0, 10.0, timestamp(0)); // Baseline

        // Alternating buy/sell of equal volume
        for i in 1..20 {
            let price = if i % 2 == 0 { 101.0 } else { 99.0 };
            let result = sampler.update(price, 10.0, timestamp(i));
            assert!(result.is_none(), "Balanced market should not emit bars");
        }

        // Imbalance should be near zero
        assert!(
            sampler.get_imbalance().abs() < 50.0,
            "Balanced market should have low imbalance"
        );
    }

    #[test]
    fn test_one_sided_flow() {
        // Strong directional flow should emit multiple bars
        let mut sampler = ImbalanceBarSampler::new(100.0, 100.0, timestamp(0));

        sampler.update(100.0, 10.0, timestamp(0)); // Baseline

        let mut bars_emitted = 0;
        for i in 1..=20 {
            let price = 100.0 + i as f64; // Monotonic price increase
            if let Some(_bar) = sampler.update(price, 30.0, timestamp(i)) {
                bars_emitted += 1;
            }
        }

        // With threshold=100 and 30 volume per tick, expect ~6 bars (600 total imbalance / 100)
        assert!(
            bars_emitted >= 5,
            "Strong directional flow should emit multiple bars, got {}",
            bars_emitted
        );
    }

    #[test]
    fn test_ewma_threshold_adaptation() {
        // EWMA threshold adapts to recent imbalance levels
        let mut sampler = ImbalanceBarSampler::new_with_ewma(100.0, 100.0, timestamp(0), 0.1);

        sampler.update(100.0, 10.0, timestamp(0)); // Baseline

        // Emit first bar
        for i in 1..=4 {
            sampler.update(100.0 + i as f64, 30.0, timestamp(i));
        }

        let initial_threshold = sampler.get_threshold();

        // Emit more bars with higher imbalance
        for i in 5..=20 {
            sampler.update(100.0 + i as f64, 50.0, timestamp(i));
        }

        let adapted_threshold = sampler.get_threshold();

        // Threshold should increase due to higher recent imbalance
        assert!(
            adapted_threshold > initial_threshold,
            "EWMA threshold should adapt upward with higher imbalance"
        );
    }

    #[test]
    fn test_multiple_bars_sequence() {
        // Test that multiple bars can be emitted in sequence
        let mut sampler = ImbalanceBarSampler::new(100.0, 50.0, timestamp(0));

        let mut bars = Vec::new();
        sampler.update(100.0, 10.0, timestamp(0));

        // Generate enough imbalance for 3 bars
        for i in 1..=10 {
            let price = 100.0 + i as f64;
            if let Some(bar) = sampler.update(price, 20.0, timestamp(i)) {
                bars.push(bar);
            }
        }

        assert!(
            bars.len() >= 3,
            "Should emit multiple bars in sequence, got {}",
            bars.len()
        );

        // Verify bars don't overlap
        for i in 1..bars.len() {
            assert!(
                bars[i].timestamp > bars[i - 1].timestamp,
                "Bars should be chronologically ordered"
            );
        }
    }

    #[test]
    fn test_zero_volume_tick() {
        // Zero volume ticks should not affect imbalance
        let mut sampler = ImbalanceBarSampler::new(100.0, 100.0, timestamp(0));

        sampler.update(100.0, 10.0, timestamp(0));
        sampler.update(101.0, 20.0, timestamp(1)); // +20 imbalance

        let imbalance_before = sampler.get_imbalance();

        sampler.update(102.0, 0.0, timestamp(2)); // Zero volume

        let imbalance_after = sampler.get_imbalance();

        assert_eq!(
            imbalance_before, imbalance_after,
            "Zero volume should not change imbalance"
        );
    }

    #[test]
    fn test_price_unchanged_tick() {
        // Price unchanged: use previous tick direction (MLFinLab convention)
        let mut sampler = ImbalanceBarSampler::new(100.0, 100.0, timestamp(0));

        sampler.update(100.0, 10.0, timestamp(0));
        sampler.update(101.0, 10.0, timestamp(1)); // Buy tick

        // Price unchanged -> should repeat last direction (buy)
        sampler.update(101.0, 10.0, timestamp(2));
        sampler.update(101.0, 10.0, timestamp(3));

        // Imbalance should continue increasing (all buy ticks)
        assert!(
            sampler.get_imbalance() >= 30.0,
            "Unchanged price should use previous direction"
        );
    }

    #[test]
    fn test_high_low_tracking() {
        // Verify high/low are correctly tracked within bar
        // Threshold=100, need cumulative imbalance of ±100 to emit bar
        let mut sampler = ImbalanceBarSampler::new(100.0, 100.0, timestamp(0));

        sampler.update(100.0, 10.0, timestamp(0)); // Open (imbalance=0, baseline)
        sampler.update(105.0, 20.0, timestamp(1)); // Buy: +20, total=+20 (high candidate)
        sampler.update(95.0, 15.0, timestamp(2)); // Sell: -15, total=+5 (low candidate)
        sampler.update(102.0, 25.0, timestamp(3)); // Buy: +25, total=+30
        sampler.update(108.0, 30.0, timestamp(4)); // Buy: +30, total=+60 (new high)
        sampler.update(103.0, 50.0, timestamp(5)); // Sell: -50, total=+10

        // Next buy tick pushes imbalance to +110 >= 100 → triggers bar
        let bar = sampler.update(104.0, 100.0, timestamp(6));
        assert!(
            bar.is_some(),
            "Bar should emit when imbalance reaches 110 (>= 100)"
        );

        let bar = bar.unwrap();
        assert_eq!(bar.open, 100.0);
        assert_eq!(bar.high, 108.0);
        assert_eq!(bar.low, 95.0);
        assert_eq!(bar.close, 104.0); // Last price before bar emission
    }
}