foxhunt/crates/ml-features/src/alternative_bars.rs

//! Alternative Bar Sampling Techniques
//!
//! Implementation of alternative bar types for improved ML model performance:
//! - Tick Bars: Aggregate every N ticks (Agent B3 - PRIMARY TASK)
//! - Volume Bars: Aggregate every N volume units
//! - Dollar Bars: Aggregate every $N traded
//! - Run Bars: Aggregate based on consecutive directional ticks
//! - Imbalance Bars: Aggregate based on buy/sell imbalance
//!
//! Based on Lopez de Prado (2018) - "Advances in Financial Machine Learning"

use chrono::{DateTime, Utc};

pub use crate::OHLCVBar;

/// Tick Bar Sampler - Aggregates every N ticks (PRIMARY IMPLEMENTATION - Agent B3)
///
/// Performance: <50μs per bar (target from Wave B Agent B3)
///
/// # Example
/// ```
/// use ml::features::alternative_bars::TickBarSampler;
/// use chrono::Utc;
///
/// let mut sampler = TickBarSampler::new(100); // 100 ticks per bar
///
/// for i in 0..150 {
///     let price = 100.0 + (i as f64 * 0.01);
///     let volume = 10.0;
///     let timestamp = Utc::now();
///
///     if let Some(bar) = sampler.update(price, volume, timestamp) {
///         println!("Bar formed: O={} H={} L={} C={} V={}",
///             bar.open, bar.high, bar.low, bar.close, bar.volume);
///     }
/// }
/// ```
#[derive(Debug)]
pub struct TickBarSampler {
    /// Number of ticks required to form a bar
    threshold: usize,
    /// Current tick count in the active bar
    tick_count: usize,
    /// Timestamp of the first tick in the current bar
    first_timestamp: Option<DateTime<Utc>>,
    /// Opening price of the current bar
    current_open: Option<f64>,
    /// Highest price seen in the current bar
    current_high: f64,
    /// Lowest price seen in the current bar
    current_low: f64,
    /// Cumulative volume in the current bar
    cumulative_volume: f64,
    /// Last price (becomes close when bar completes)
    last_price: f64,
}

impl TickBarSampler {
    /// Create a new tick bar sampler
    ///
    /// # Arguments
    /// * `threshold` - Number of ticks per bar (e.g., 100, 1000)
    ///
    /// # Panics
    /// Panics if threshold is 0
    pub fn new(threshold: usize) -> Self {
        assert!(threshold > 0, "Threshold must be greater than 0");

        Self {
            threshold,
            tick_count: 0,
            first_timestamp: None,
            current_open: None,
            current_high: f64::NEG_INFINITY,
            current_low: f64::INFINITY,
            cumulative_volume: 0.0,
            last_price: 0.0,
        }
    }

    /// Process a single tick and return completed bar if threshold reached
    ///
    /// # Arguments
    /// * `price` - Trade price
    /// * `volume` - Trade volume (can be 0)
    /// * `timestamp` - Trade timestamp
    ///
    /// # Returns
    /// `Some(OHLCVBar)` if a bar was completed, `None` otherwise
    pub fn update(
        &mut self,
        price: f64,
        volume: f64,
        timestamp: DateTime<Utc>,
    ) -> Option<OHLCVBar> {
        // Initialize on first tick
        if self.current_open.is_none() {
            self.current_open = Some(price);
            self.first_timestamp = Some(timestamp);
        }

        // Update OHLCV
        self.current_high = self.current_high.max(price);
        self.current_low = self.current_low.min(price);
        self.cumulative_volume += volume;
        self.last_price = price;

        // Increment tick count
        self.tick_count += 1;

        // Check if bar is complete
        (self.tick_count >= self.threshold).then(|| {
            let bar = OHLCVBar {
                timestamp: self.first_timestamp.unwrap_or(timestamp),
                open: self.current_open.unwrap_or(price),
                high: self.current_high,
                low: self.current_low,
                close: self.last_price,
                volume: self.cumulative_volume,
            };

            // Reset for next bar
            self.reset();

            bar
        })
    }

    /// Get the tick threshold
    pub const fn threshold(&self) -> usize {
        self.threshold
    }

    /// Get the current tick count (0 to threshold-1)
    pub const fn tick_count(&self) -> usize {
        self.tick_count
    }

    /// Reset the sampler state for a new bar
    const fn reset(&mut self) {
        self.tick_count = 0;
        self.first_timestamp = None;
        self.current_open = None;
        self.current_high = f64::NEG_INFINITY;
        self.current_low = f64::INFINITY;
        self.cumulative_volume = 0.0;
    }
}

// Additional samplers for future agents

/// Volume Bar Sampler - Aggregates every N volume units
#[derive(Debug)]
pub struct VolumeBarSampler {
    threshold: u64,
    cumulative_volume: u64,
    first_timestamp: Option<DateTime<Utc>>,
    current_open: Option<f64>,
    current_high: f64,
    current_low: f64,
    last_price: f64,
}

impl VolumeBarSampler {
    pub fn new(threshold: u64) -> Self {
        assert!(threshold > 0, "Threshold must be greater than 0");
        Self {
            threshold,
            cumulative_volume: 0,
            first_timestamp: None,
            current_open: None,
            current_high: f64::NEG_INFINITY,
            current_low: f64::INFINITY,
            last_price: 0.0,
        }
    }

    pub fn update(
        &mut self,
        price: f64,
        volume: f64,
        timestamp: DateTime<Utc>,
    ) -> Option<OHLCVBar> {
        let volume_units = volume.round() as u64;
        self.cumulative_volume += volume_units;

        if self.current_open.is_none() {
            self.current_open = Some(price);
            self.first_timestamp = Some(timestamp);
        }

        self.current_high = self.current_high.max(price);
        self.current_low = self.current_low.min(price);
        self.last_price = price;

        (self.cumulative_volume >= self.threshold).then(|| {
            let bar = OHLCVBar {
                timestamp: self.first_timestamp.unwrap_or(timestamp),
                open: self.current_open.unwrap_or(price),
                high: self.current_high,
                low: self.current_low,
                close: self.last_price,
                volume: self.cumulative_volume as f64,
            };
            self.reset();
            bar
        })
    }

    pub const fn threshold(&self) -> u64 {
        self.threshold
    }

    pub const fn cumulative_volume(&self) -> u64 {
        self.cumulative_volume
    }

    const fn reset(&mut self) {
        self.cumulative_volume = 0;
        self.first_timestamp = None;
        self.current_open = None;
        self.current_high = f64::NEG_INFINITY;
        self.current_low = f64::INFINITY;
    }
}

/// Dollar Bar Sampler - Aggregates every $N traded
#[derive(Debug)]
pub struct DollarBarSampler {
    threshold: f64,
    cumulative_dollar: f64,
    first_timestamp: Option<DateTime<Utc>>,
    current_open: Option<f64>,
    current_high: f64,
    current_low: f64,
    cumulative_volume: f64,
    last_price: f64,
    adaptive_mode: bool,
    ewma_alpha: f64,
}

impl DollarBarSampler {
    pub fn new(threshold: f64) -> Self {
        assert!(threshold > 0.0, "Threshold must be greater than 0");
        Self {
            threshold,
            cumulative_dollar: 0.0,
            first_timestamp: None,
            current_open: None,
            current_high: f64::NEG_INFINITY,
            current_low: f64::INFINITY,
            cumulative_volume: 0.0,
            last_price: 0.0,
            adaptive_mode: false,
            ewma_alpha: 0.0,
        }
    }

    /// Create adaptive dollar bar sampler with EWMA threshold adjustment
    pub fn new_adaptive(initial_threshold: f64, alpha: f64) -> Self {
        assert!(
            initial_threshold > 0.0,
            "Initial threshold must be positive"
        );
        assert!(alpha > 0.0 && alpha <= 1.0, "Alpha must be in (0, 1]");
        Self {
            threshold: initial_threshold,
            cumulative_dollar: 0.0,
            first_timestamp: None,
            current_open: None,
            current_high: f64::NEG_INFINITY,
            current_low: f64::INFINITY,
            cumulative_volume: 0.0,
            last_price: 0.0,
            adaptive_mode: true,
            ewma_alpha: alpha,
        }
    }

    /// Get current threshold (for test compatibility)
    pub const fn get_threshold(&self) -> f64 {
        self.threshold
    }

    pub fn update(
        &mut self,
        price: f64,
        volume: f64,
        timestamp: DateTime<Utc>,
    ) -> Option<OHLCVBar> {
        // Validate inputs
        assert!(price >= 0.0, "Price cannot be negative");
        assert!(volume >= 0.0, "Volume cannot be negative");

        // Ignore zero-volume ticks
        if volume == 0.0 {
            return None;
        }

        let dollar_value = price * volume;
        self.cumulative_dollar += dollar_value;

        if self.current_open.is_none() {
            self.current_open = Some(price);
            self.first_timestamp = Some(timestamp);
        }

        self.current_high = self.current_high.max(price);
        self.current_low = self.current_low.min(price);
        self.cumulative_volume += volume;
        self.last_price = price;

        (self.cumulative_dollar >= self.threshold).then(|| {
            let bar = OHLCVBar {
                timestamp: self.first_timestamp.unwrap_or(timestamp),
                open: self.current_open.unwrap_or(price),
                high: self.current_high,
                low: self.current_low,
                close: self.last_price,
                volume: self.cumulative_volume,
            };

            // Update threshold if adaptive mode (EWMA)
            if self.adaptive_mode {
                self.threshold = self.ewma_alpha * self.threshold
                    + (1.0 - self.ewma_alpha) * self.cumulative_dollar;
            }

            self.reset();
            bar
        })
    }

    pub const fn threshold(&self) -> f64 {
        self.threshold
    }

    pub const fn cumulative_dollar(&self) -> f64 {
        self.cumulative_dollar
    }

    /// Get accumulated dollar volume (test compatibility alias)
    pub const fn get_accumulated(&self) -> f64 {
        self.cumulative_dollar
    }

    const fn reset(&mut self) {
        self.cumulative_dollar = 0.0;
        self.first_timestamp = None;
        self.current_open = None;
        self.current_high = f64::NEG_INFINITY;
        self.current_low = f64::INFINITY;
        self.cumulative_volume = 0.0;
    }
}

/// Imbalance Bar Sampler - Aggregates based on buy/sell imbalance
///
/// Emits bars when cumulative imbalance exceeds threshold.
/// Buy ticks (price increase) add to imbalance, sell ticks (price decrease) subtract.
///
/// Based on Lopez de Prado (2018) - "Advances in Financial Machine Learning"
///
/// Performance: <50μs per bar (Wave B target)
///
/// # Example
/// ```
/// use ml::features::alternative_bars::ImbalanceBarSampler;
/// use chrono::Utc;
///
/// let mut sampler = ImbalanceBarSampler::new(100.0, 100.0, Utc::now());
///
/// // Process ticks
/// sampler.update(100.0, 10.0, Utc::now()); // Baseline
/// sampler.update(101.0, 20.0, Utc::now()); // Buy tick: +20 imbalance
/// sampler.update(102.0, 30.0, Utc::now()); // Buy tick: +30 imbalance
///
/// // Bar emits when imbalance >= 100
/// if let Some(bar) = sampler.update(103.0, 60.0, Utc::now()) {
///     println!("Bar: O={} H={} L={} C={} V={}", bar.open, bar.high, bar.low, bar.close, bar.volume);
/// }
/// ```
#[derive(Debug)]
pub struct ImbalanceBarSampler {
    /// Imbalance threshold for bar formation
    threshold: f64,
    /// Cumulative buy/sell imbalance (positive=buy, negative=sell)
    cumulative_imbalance: f64,
    /// Previous price for tick direction classification
    previous_price: Option<f64>,
    /// Last tick direction (+1=buy, -1=sell, 0=unchanged)
    last_direction: i8,
    /// OHLCV tracking
    first_timestamp: Option<DateTime<Utc>>,
    current_open: Option<f64>,
    current_high: f64,
    current_low: f64,
    cumulative_volume: f64,
    last_price: f64,
    /// EWMA threshold adaptation
    adaptive_mode: bool,
    ewma_alpha: f64,
}

impl ImbalanceBarSampler {
    /// Create a new imbalance bar sampler with fixed threshold
    ///
    /// # Arguments
    /// * `initial_price` - Starting price (for direction classification)
    /// * `threshold` - Imbalance threshold (e.g., 100.0 for ±100 units)
    /// * `timestamp` - Initial timestamp
    pub fn new(initial_price: f64, threshold: f64, timestamp: DateTime<Utc>) -> Self {
        assert!(threshold > 0.0, "Threshold must be greater than 0");

        Self {
            threshold,
            cumulative_imbalance: 0.0,
            previous_price: Some(initial_price),
            last_direction: 0,
            first_timestamp: Some(timestamp),
            current_open: None,
            current_high: f64::NEG_INFINITY,
            current_low: f64::INFINITY,
            cumulative_volume: 0.0,
            last_price: initial_price,
            adaptive_mode: false,
            ewma_alpha: 0.0,
        }
    }

    /// Create adaptive imbalance bar sampler with EWMA threshold adjustment
    ///
    /// # Arguments
    /// * `initial_price` - Starting price
    /// * `threshold` - Initial imbalance threshold
    /// * `timestamp` - Initial timestamp
    /// * `alpha` - EWMA smoothing factor (0 < alpha <= 1, e.g., 0.1)
    pub fn new_with_ewma(
        initial_price: f64,
        threshold: f64,
        timestamp: DateTime<Utc>,
        alpha: f64,
    ) -> Self {
        assert!(threshold > 0.0, "Threshold must be greater than 0");
        assert!(alpha > 0.0 && alpha <= 1.0, "Alpha must be in (0, 1]");

        Self {
            threshold,
            cumulative_imbalance: 0.0,
            previous_price: Some(initial_price),
            last_direction: 0,
            first_timestamp: Some(timestamp),
            current_open: None,
            current_high: f64::NEG_INFINITY,
            current_low: f64::INFINITY,
            cumulative_volume: 0.0,
            last_price: initial_price,
            adaptive_mode: true,
            ewma_alpha: alpha,
        }
    }

    /// Process a tick and return completed bar if threshold exceeded
    ///
    /// # Arguments
    /// * `price` - Trade price
    /// * `volume` - Trade volume
    /// * `timestamp` - Trade timestamp
    ///
    /// # Returns
    /// `Some(OHLCVBar)` if imbalance threshold was exceeded, `None` otherwise
    ///
    /// # Tick Classification
    /// - Buy tick: `price > previous_price` → direction = +1
    /// - Sell tick: `price < previous_price` → direction = -1
    /// - Unchanged: `price == previous_price` → use `last_direction` (`MLFinLab` convention)
    pub fn update(
        &mut self,
        price: f64,
        volume: f64,
        timestamp: DateTime<Utc>,
    ) -> Option<OHLCVBar> {
        // Ignore zero-volume ticks
        if volume == 0.0 {
            return None;
        }

        // Initialize on first tick
        if self.current_open.is_none() {
            self.current_open = Some(price);
            self.first_timestamp = Some(timestamp);
        }

        // Classify tick direction
        let direction = if let Some(prev_price) = self.previous_price {
            if price > prev_price {
                1 // Buy tick
            } else if price < prev_price {
                -1 // Sell tick
            } else {
                // Price unchanged: use previous direction (MLFinLab convention)
                self.last_direction
            }
        } else {
            0 // First tick has no direction
        };

        // Update imbalance: positive for buys, negative for sells
        let imbalance_contribution = direction as f64 * volume;
        self.cumulative_imbalance += imbalance_contribution;

        // Update OHLCV
        self.current_high = self.current_high.max(price);
        self.current_low = self.current_low.min(price);
        self.cumulative_volume += volume;
        self.last_price = price;

        // Update state for next tick
        self.previous_price = Some(price);
        self.last_direction = direction;

        // Check if bar should be emitted (absolute imbalance >= threshold)
        (self.cumulative_imbalance.abs() >= self.threshold).then(|| {
            let bar = OHLCVBar {
                timestamp: self.first_timestamp.unwrap_or(timestamp),
                open: self.current_open.unwrap_or(price),
                high: self.current_high,
                low: self.current_low,
                close: self.last_price,
                volume: self.cumulative_volume,
            };

            // Update threshold if adaptive mode (EWMA)
            if self.adaptive_mode {
                let observed_imbalance = self.cumulative_imbalance.abs();
                self.threshold =
                    self.ewma_alpha * self.threshold + (1.0 - self.ewma_alpha) * observed_imbalance;
            }

            // Reset for next bar
            self.reset();

            bar
        })
    }

    /// Get current cumulative imbalance
    pub const fn get_imbalance(&self) -> f64 {
        self.cumulative_imbalance
    }

    /// Get current threshold
    pub const fn get_threshold(&self) -> f64 {
        self.threshold
    }

    /// Reset the sampler state for a new bar
    const fn reset(&mut self) {
        self.cumulative_imbalance = 0.0;
        self.first_timestamp = None;
        self.current_open = None;
        self.current_high = f64::NEG_INFINITY;
        self.current_low = f64::INFINITY;
        self.cumulative_volume = 0.0;
        // Keep previous_price and last_direction for continuity
    }
}

/// Run Bar Sampler - Aggregates based on consecutive directional ticks
///
/// Emits bars when consecutive buy or sell ticks exceed threshold.
/// A "run" is a sequence of ticks moving in the same direction.
///
/// Based on Lopez de Prado (2018) - "Advances in Financial Machine Learning"
///
/// Performance: <50μs per bar (Wave B target)
///
/// # Example
/// ```
/// use ml::features::alternative_bars::RunBarSampler;
/// use chrono::Utc;
///
/// let mut sampler = RunBarSampler::new(5); // 5 consecutive ticks in same direction
///
/// // Send 5 buy ticks (price increasing)
/// sampler.update(100.0, 10.0, Utc::now());
/// sampler.update(100.1, 10.0, Utc::now());
/// sampler.update(100.2, 10.0, Utc::now());
/// sampler.update(100.3, 10.0, Utc::now());
///
/// // 5th buy tick triggers bar
/// if let Some(bar) = sampler.update(100.4, 10.0, Utc::now()) {
///     println!("Bar: O={} H={} L={} C={} V={}", bar.open, bar.high, bar.low, bar.close, bar.volume);
/// }
/// ```
#[derive(Debug)]
pub struct RunBarSampler {
    /// Threshold for consecutive directional ticks
    threshold: usize,
    /// Current run count (consecutive ticks in same direction)
    run_count: usize,
    /// Previous price for direction classification
    previous_price: Option<f64>,
    /// Current direction (+1=buy, -1=sell, 0=no direction)
    current_direction: i8,
    /// OHLCV tracking
    first_timestamp: Option<DateTime<Utc>>,
    current_open: Option<f64>,
    current_high: f64,
    current_low: f64,
    cumulative_volume: f64,
    last_price: f64,
}

impl RunBarSampler {
    /// Create a new run bar sampler
    ///
    /// # Arguments
    /// * `threshold` - Number of consecutive directional ticks per bar (e.g., 5, 10)
    ///
    /// # Panics
    /// Panics if threshold is 0
    pub fn new(threshold: usize) -> Self {
        assert!(threshold > 0, "Threshold must be greater than 0");

        Self {
            threshold,
            run_count: 0,
            previous_price: None,
            current_direction: 0,
            first_timestamp: None,
            current_open: None,
            current_high: f64::NEG_INFINITY,
            current_low: f64::INFINITY,
            cumulative_volume: 0.0,
            last_price: 0.0,
        }
    }

    /// Process a tick and return completed bar if run threshold reached
    ///
    /// # Arguments
    /// * `price` - Trade price
    /// * `volume` - Trade volume
    /// * `timestamp` - Trade timestamp
    ///
    /// # Returns
    /// `Some(OHLCVBar)` if consecutive run threshold was reached, `None` otherwise
    ///
    /// # Direction Classification
    /// - Buy tick: `price > previous_price` → direction = +1
    /// - Sell tick: `price < previous_price` → direction = -1
    /// - Unchanged: `price == previous_price` → no direction (run continues)
    /// - Direction change: Resets `run_count` to 1
    pub fn update(
        &mut self,
        price: f64,
        volume: f64,
        timestamp: DateTime<Utc>,
    ) -> Option<OHLCVBar> {
        // Determine tick direction FIRST (before updating state)
        let direction = if let Some(prev_price) = self.previous_price {
            if price > prev_price {
                1 // Buy tick
            } else if price < prev_price {
                -1 // Sell tick
            } else {
                0 // No direction (price unchanged)
            }
        } else {
            0 // First tick has no direction
        };

        // Direction change detection: if we have a new direction (not 0) different from current
        let direction_changed =
            direction != 0 && self.current_direction != 0 && direction != self.current_direction;

        if direction_changed {
            // Direction changed - emit bar if threshold was met in previous run
            if self.run_count >= self.threshold {
                let bar = OHLCVBar {
                    timestamp: self.first_timestamp.unwrap_or(timestamp),
                    open: self.current_open.unwrap_or(price),
                    high: self.current_high,
                    low: self.current_low,
                    close: self.last_price,
                    volume: self.cumulative_volume,
                };

                // Reset and start new run with current tick
                self.reset();
                self.current_open = Some(price);
                self.first_timestamp = Some(timestamp);
                self.current_high = price;
                self.current_low = price;
                self.cumulative_volume = volume;
                self.last_price = price;
                self.run_count = 1;
                self.current_direction = direction;
                self.previous_price = Some(price);

                return Some(bar);
            } else {
                // Direction changed but threshold not met - reset and start new run
                self.reset();
                self.current_open = Some(price);
                self.first_timestamp = Some(timestamp);
                self.current_high = price;
                self.current_low = price;
                self.cumulative_volume = volume;
                self.last_price = price;
                self.run_count = 1;
                self.current_direction = direction;
                self.previous_price = Some(price);

                return None;
            }
        }

        // No direction change - continue accumulating
        // Initialize on first tick
        if self.current_open.is_none() {
            self.current_open = Some(price);
            self.first_timestamp = Some(timestamp);
        }

        // Update OHLCV accumulation
        self.current_high = self.current_high.max(price);
        self.current_low = self.current_low.min(price);
        self.cumulative_volume += volume;
        self.last_price = price;

        // Update previous price for next comparison
        self.previous_price = Some(price);

        // Increment run count on EVERY tick
        self.run_count += 1;

        // Set direction on first directional tick
        if direction != 0 && self.current_direction == 0 {
            self.current_direction = direction;
        }

        // Check if threshold reached AND we have a direction
        (self.run_count >= self.threshold && self.current_direction != 0).then(|| {
            let bar = OHLCVBar {
                timestamp: self.first_timestamp.unwrap_or(timestamp),
                open: self.current_open.unwrap_or(price),
                high: self.current_high,
                low: self.current_low,
                close: self.last_price,
                volume: self.cumulative_volume,
            };

            // Reset for next bar
            self.reset();

            bar
        })
    }

    /// Get the run threshold
    pub const fn threshold(&self) -> usize {
        self.threshold
    }

    /// Get the current run count
    pub const fn run_count(&self) -> usize {
        self.run_count
    }

    /// Get the current direction (-1 for sell, 0 for neutral, +1 for buy)
    pub const fn direction(&self) -> i8 {
        self.current_direction
    }

    /// Reset the sampler state for a new bar
    pub const fn reset(&mut self) {
        self.run_count = 0;
        self.current_direction = 0;
        self.first_timestamp = None;
        self.current_open = None;
        self.current_high = f64::NEG_INFINITY;
        self.current_low = f64::INFINITY;
        self.cumulative_volume = 0.0;
        // Keep previous_price for continuity
    }
}