foxhunt/crates/ml-labeling/src/triple_barrier.rs

//! Triple Barrier Engine implementation
//!
//! High-performance triple barrier labeling with <80us latency target.
//! Based on the Python reference implementation from `HFTTrendfollowing`
//! with optimizations for ultra-low latency financial applications.

use std::collections::VecDeque;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Instant;

use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use super::constants::BASIS_POINTS_PER_DOLLAR;
use super::types::{BarrierConfig, BarrierTouchedFirst, BarrierResult, EventLabel};

/// Price point for tracking
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct PricePoint {
    pub price_cents: u64,
    pub timestamp_ns: u64,
}

impl PricePoint {
    pub const fn new(price_cents: u64, timestamp_ns: u64) -> Self {
        Self {
            price_cents,
            timestamp_ns,
        }
    }
}

/// Triple barrier tracker for a single position
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BarrierTracker {
    pub entry_price_cents: u64,
    pub entry_timestamp_ns: u64,
    pub upper_barrier_cents: u64,
    pub lower_barrier_cents: u64,
    pub expiry_timestamp_ns: u64,
    pub config: BarrierConfig,
    pub touched_first: Option<BarrierTouchedFirst>,
    pub final_result: Option<BarrierResult>,
}

impl BarrierTracker {
    pub const fn new(entry_price_cents: u64, entry_timestamp_ns: u64, config: BarrierConfig) -> Self {
        let upper_barrier_cents = entry_price_cents
            + (entry_price_cents * config.profit_target_bps as u64)
                / BASIS_POINTS_PER_DOLLAR as u64;
        let lower_barrier_cents = entry_price_cents
            - (entry_price_cents * config.stop_loss_bps as u64) / BASIS_POINTS_PER_DOLLAR as u64;
        let expiry_timestamp_ns = entry_timestamp_ns + config.max_holding_period_ns;

        Self {
            entry_price_cents,
            entry_timestamp_ns,
            upper_barrier_cents,
            lower_barrier_cents,
            expiry_timestamp_ns,
            config,
            touched_first: None,
            final_result: None,
        }
    }

    /// Update tracker with new price data
    pub fn update(&mut self, price_point: PricePoint) -> Option<EventLabel> {
        if self.final_result.is_some() {
            return None; // Already closed
        }

        // Check if expired
        if price_point.timestamp_ns >= self.expiry_timestamp_ns {
            self.final_result = Some(BarrierResult::TimeExpiry);
            return Some(self.create_event_label(price_point));
        }

        // Check barriers
        if price_point.price_cents >= self.upper_barrier_cents {
            if self.touched_first.is_none() {
                self.touched_first = Some(BarrierTouchedFirst::Upper);
            }
            self.final_result = Some(BarrierResult::ProfitTarget);
            return Some(self.create_event_label(price_point));
        }

        if price_point.price_cents <= self.lower_barrier_cents {
            if self.touched_first.is_none() {
                self.touched_first = Some(BarrierTouchedFirst::Lower);
            }
            self.final_result = Some(BarrierResult::StopLoss);
            return Some(self.create_event_label(price_point));
        }

        None
    }

    fn create_event_label(&self, price_point: PricePoint) -> EventLabel {
        let return_bps = if self.entry_price_cents == 0 {
            0 // Guard against divide-by-zero when entry price is unknown
        } else if price_point.price_cents > self.entry_price_cents {
            (((price_point.price_cents - self.entry_price_cents) as i64 * BASIS_POINTS_PER_DOLLAR)
                / self.entry_price_cents as i64) as i32
        } else {
            (-((self.entry_price_cents - price_point.price_cents) as i64 * BASIS_POINTS_PER_DOLLAR)
                / self.entry_price_cents as i64) as i32
        };

        let label_value = match self.final_result {
            Some(BarrierResult::ProfitTarget) => 1,
            Some(BarrierResult::StopLoss) => -1,
            Some(BarrierResult::TimeExpiry) => {
                if return_bps > 0 {
                    1
                } else if return_bps < 0 {
                    -1
                } else {
                    0
                }
            },
            None => 0,
        };

        EventLabel {
            event_timestamp_ns: price_point.timestamp_ns,
            entry_price_cents: self.entry_price_cents,
            barrier_result: self.final_result.unwrap_or(BarrierResult::TimeExpiry),
            label_value,
            return_bps,
            quality_score: self.calculate_quality_score(price_point),
            processing_latency_us: 0, // Will be filled by engine
        }
    }

    const fn calculate_quality_score(&self, _price_point: PricePoint) -> f64 {
        // Simple quality score based on how quickly the barrier was hit
        match self.final_result {
            Some(BarrierResult::ProfitTarget) => 0.9,
            Some(BarrierResult::StopLoss) => 0.8,
            Some(BarrierResult::TimeExpiry) => 0.5,
            None => 0.0,
        }
    }

    pub const fn is_closed(&self) -> bool {
        self.final_result.is_some()
    }
}

/// High-performance triple barrier labeling engine
#[derive(Debug)]
pub struct TripleBarrierEngine {
    active_trackers: DashMap<Uuid, BarrierTracker>,
    completed_labels: VecDeque<EventLabel>,
    stats: Arc<AtomicU64>,
    max_active_trackers: usize,
}

impl TripleBarrierEngine {
    pub fn new(max_active_trackers: usize) -> Self {
        Self {
            active_trackers: DashMap::new(),
            completed_labels: VecDeque::new(),
            stats: Arc::new(AtomicU64::new(0)),
            max_active_trackers,
        }
    }

    /// Start tracking a new position
    pub fn start_tracking(
        &mut self,
        config: BarrierConfig,
        entry_price_cents: u64,
        entry_timestamp_ns: u64,
    ) -> Result<Uuid, String> {
        if self.active_trackers.len() >= self.max_active_trackers {
            return Err("Maximum active trackers reached".to_owned());
        }

        let tracker_id = Uuid::new_v4();
        let tracker = BarrierTracker::new(entry_price_cents, entry_timestamp_ns, config);

        self.active_trackers.insert(tracker_id, tracker);
        Ok(tracker_id)
    }

    /// Update all trackers with new price data
    pub fn update_all(&mut self, price_point: PricePoint) -> Vec<EventLabel> {
        let start = Instant::now();
        let mut completed_labels = Vec::new();
        let mut trackers_to_remove = Vec::new();

        for mut entry in self.active_trackers.iter_mut() {
            let tracker_id = *entry.key();
            let tracker = entry.value_mut();

            if let Some(mut label) = tracker.update(price_point) {
                label.processing_latency_us = start.elapsed().as_micros() as u32;
                completed_labels.push(label);
                trackers_to_remove.push(tracker_id);
            }
        }

        // Remove completed trackers
        for tracker_id in trackers_to_remove {
            self.active_trackers.remove(&tracker_id);
        }

        // Store completed labels
        for label in &completed_labels {
            self.completed_labels.push_back(label.clone());
        }

        // Update stats
        self.stats
            .fetch_add(completed_labels.len() as u64, Ordering::Relaxed);

        completed_labels
    }

    /// Update specific tracker
    pub fn update_tracker(
        &mut self,
        tracker_id: Uuid,
        price_point: PricePoint,
    ) -> Option<EventLabel> {
        let start = Instant::now();

        if let Some(mut entry) = self.active_trackers.get_mut(&tracker_id) {
            let tracker = entry.value_mut();

            if let Some(mut label) = tracker.update(price_point) {
                label.processing_latency_us = start.elapsed().as_micros() as u32;
                self.completed_labels.push_back(label.clone());
                self.stats.fetch_add(1, Ordering::Relaxed);

                // Remove completed tracker
                drop(entry);
                self.active_trackers.remove(&tracker_id);

                return Some(label);
            }
        }

        None
    }

    /// Get completed labels and clear the buffer
    pub fn drain_completed_labels(&mut self) -> Vec<EventLabel> {
        self.completed_labels.drain(..).collect()
    }

    /// Get number of active trackers
    pub fn active_count(&self) -> usize {
        self.active_trackers.len()
    }

    /// Get total completed labels
    pub fn completed_count(&self) -> u64 {
        self.stats.load(Ordering::Relaxed)
    }

    /// Force expire old trackers
    pub fn expire_old_trackers(&mut self, current_timestamp_ns: u64) -> Vec<EventLabel> {
        let start = Instant::now();
        let mut expired_labels = Vec::new();
        let mut trackers_to_remove = Vec::new();

        for entry in &self.active_trackers {
            let tracker_id = *entry.key();
            let tracker = entry.value();

            if current_timestamp_ns >= tracker.expiry_timestamp_ns {
                let price_point = PricePoint::new(tracker.entry_price_cents, current_timestamp_ns);
                let mut tracker_clone = tracker.clone();

                if let Some(mut label) = tracker_clone.update(price_point) {
                    label.processing_latency_us = start.elapsed().as_micros() as u32;
                    expired_labels.push(label);
                    trackers_to_remove.push(tracker_id);
                }
            }
        }

        // Remove expired trackers
        for tracker_id in trackers_to_remove {
            self.active_trackers.remove(&tracker_id);
        }

        // Store expired labels
        for label in &expired_labels {
            self.completed_labels.push_back(label.clone());
        }

        self.stats
            .fetch_add(expired_labels.len() as u64, Ordering::Relaxed);
        expired_labels
    }

    /// Get tracker by ID
    pub fn get_tracker(&self, tracker_id: &Uuid) -> Option<BarrierTracker> {
        self.active_trackers
            .get(tracker_id)
            .map(|entry| entry.value().clone())
    }

    /// Clear all trackers (for testing)
    pub fn clear(&mut self) {
        self.active_trackers.clear();
        self.completed_labels.clear();
        self.stats.store(0, Ordering::Relaxed);
    }
}

#[cfg(test)]
#[allow(
    clippy::inconsistent_digit_grouping,
    clippy::unnecessary_wraps,
    clippy::let_underscore_must_use,
    clippy::len_zero
)]
mod tests {
    use super::*;

    #[test]
    fn test_barrier_tracker_creation() {
        let config = BarrierConfig::conservative();
        let entry_price_cents = 10000; // $100.00
        let entry_timestamp_ns = 1692000000_000_000_000;

        let tracker = BarrierTracker::new(entry_price_cents, entry_timestamp_ns, config);

        assert_eq!(tracker.entry_price_cents, 10000);
        assert_eq!(tracker.upper_barrier_cents, 10100); // +1%
        assert_eq!(tracker.lower_barrier_cents, 9950); // -0.5%
    }

    #[test]
    fn test_barrier_touching() -> Result<(), Box<dyn std::error::Error>> {
        let config = BarrierConfig::conservative();
        let mut tracker = BarrierTracker::new(10000, 1692000000_000_000_000, config);

        // Test profit barrier hit
        let profit_price = PricePoint::new(10150, 1692000000_000_000_000 + 1800_000_000_000);
        let result = tracker.update(profit_price);

        assert!(result.is_some());
        let label = result.ok_or("Expected label from tracker update")?;
        assert_eq!(label.label_value, 1);
        assert!(label.return_bps > 0);
        assert!(matches!(label.barrier_result, BarrierResult::ProfitTarget));

        Ok(())
    }

    #[test]
    fn test_engine_creation() {
        let engine = TripleBarrierEngine::new(1000);
        assert_eq!(engine.active_count(), 0);
        assert_eq!(engine.completed_count(), 0);
    }

    #[test]
    fn test_engine_tracking() -> Result<(), Box<dyn std::error::Error>> {
        let mut engine = TripleBarrierEngine::new(1000);
        let config = BarrierConfig::conservative();

        let _tracker_id = engine.start_tracking(config, 10000, 1692000000_000_000_000);
        assert_eq!(engine.active_count(), 1);

        // Update with profit-taking price
        let price_point = PricePoint::new(10150, 1692000000_000_000_000 + 1000_000_000);
        let labels = engine.update_all(price_point);

        assert_eq!(labels.len(), 1);
        assert_eq!(engine.active_count(), 0);
        assert_eq!(engine.completed_count(), 1);

        Ok(())
    }

    #[test]
    fn test_time_expiry() -> Result<(), Box<dyn std::error::Error>> {
        let mut engine = TripleBarrierEngine::new(1000);
        let config = BarrierConfig::conservative();

        let _tracker_id = engine.start_tracking(config, 10000, 1692000000_000_000_000);

        // Force expire
        let expired_labels = engine.expire_old_trackers(1692000000_000_000_000 + 3700_000_000_000); // 1 hour + 100 seconds

        assert_eq!(expired_labels.len(), 1);
        assert!(matches!(
            expired_labels[0].barrier_result,
            BarrierResult::TimeExpiry
        ));
        assert_eq!(engine.active_count(), 0);

        Ok(())
    }

    #[test]
    fn test_quality_score_calculation() -> Result<(), Box<dyn std::error::Error>> {
        let config = BarrierConfig::conservative();
        let mut tracker = BarrierTracker::new(10000, 1692000000_000_000_000, config);

        let profit_price = PricePoint::new(10150, 1692000000_000_000_000 + 1000_000_000);
        let result = tracker.update(profit_price);

        assert!(result.is_some());
        let label = result.ok_or("Expected label from tracker update")?;
        assert!(label.quality_score > 0.8); // Profit targets should have high quality

        Ok(())
    }

    #[test]
    fn test_multiple_updates() {
        let mut engine = TripleBarrierEngine::new(1000);
        let config = BarrierConfig::conservative();

        // Start multiple trackers
        for i in 0..5 {
            let _ = engine.start_tracking(config.clone(), 10000 + i * 100, 1692000000_000_000_000);
        }

        assert_eq!(engine.active_count(), 5);

        // Update with various prices
        let price_point = PricePoint::new(10200, 1692000000_000_000_000 + 1000_000_000);
        let labels = engine.update_all(price_point);

        // Some should hit profit target
        assert!(labels.len() > 0);
        assert!(engine.active_count() < 5);
    }
}