foxhunt/crates/ml/src/labeling/sample_weights.rs

//! Sample weighting algorithms for training data enhancement
//!
//! Implements volatility/return/time-based weighting to improve ML model training.

use serde::{Deserialize, Serialize};

use super::gpu_acceleration::LabelingError;
use super::types::{EventLabel, WeightedSample};

/// Configuration for sample weighting calculation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WeightingConfig {
    /// Time decay factor for recency weighting
    pub time_decay: f64,
    /// Return scaling factor for return-based weighting
    pub return_scale: f64,
    /// Volatility scaling factor for volatility-based weighting
    pub volatility_scale: f64,
}

impl WeightingConfig {
    /// Standard configuration values
    pub fn standard() -> Self {
        Self {
            time_decay: 0.95,
            return_scale: 1.0,
            volatility_scale: 1.0,
        }
    }
}

/// Calculator for sample weights
#[derive(Debug)]
pub struct SampleWeightCalculator {
    config: WeightingConfig,
}

impl SampleWeightCalculator {
    /// Create new calculator with given configuration
    pub fn new(config: WeightingConfig) -> Self {
        Self { config }
    }

    /// Calculate weights for given labels
    pub fn calculate_weights(
        &self,
        labels: &[EventLabel],
    ) -> Result<Vec<WeightedSample>, LabelingError> {
        if labels.is_empty() {
            return Ok(Vec::new());
        }

        let mut samples = Vec::with_capacity(labels.len());

        for label in labels {
            let time_weight = self.calculate_time_weight(label.event_timestamp_ns as i64, labels);
            let return_weight = self.calculate_return_weight(label.return_bps);
            let volatility_weight = self.calculate_volatility_weight(0.2); // Default volatility for now

            let combined_weight = time_weight * return_weight * volatility_weight;

            // Create features vector from the label data
            let features = vec![
                label.entry_price_cents as f64 / 100.0, // Price in dollars
                label.return_as_ratio(),                // Return as ratio
                label.quality_score,                    // Quality score
            ];

            samples.push(WeightedSample {
                timestamp_ns: label.event_timestamp_ns,
                features,
                label: label.label_value,
                weight: combined_weight,
                sample_id: None,
            });
        }

        Ok(samples)
    }

    fn calculate_time_weight(&self, timestamp_ns: i64, all_labels: &[EventLabel]) -> f64 {
        if all_labels.is_empty() {
            return 1.0;
        }

        let latest_time = all_labels
            .iter()
            .map(|l| l.event_timestamp_ns as i64)
            .max()
            .unwrap_or(timestamp_ns); // Default to current timestamp if no labels
        let time_diff_hours = (latest_time - timestamp_ns) as f64 / 3_600_000_000_000.0;
        self.config.time_decay.powf(time_diff_hours.max(0.0))
    }

    fn calculate_return_weight(&self, return_bps: i32) -> f64 {
        (return_bps.abs() as f64 / 100.0 * self.config.return_scale).max(0.1)
    }

    fn calculate_volatility_weight(&self, volatility: f64) -> f64 {
        (volatility * self.config.volatility_scale).max(0.1)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::labeling::types::BarrierResult;

    #[test]
    fn test_sample_weight_calculator() -> Result<(), LabelingError> {
        let config = WeightingConfig::standard();
        let calculator = SampleWeightCalculator::new(config);

        // Create test labels with varying returns
        let mut labels = Vec::new();
        let returns = [100, 200, 50, 300, 150]; // Basis points

        for (i, return_bps) in returns.into_iter().enumerate() {
            let barrier_result = BarrierResult::ProfitTarget;

            let label = EventLabel::new(
                (1692000000_000_000_000 + i as u64 * 3600_000_000_000)
                    .saturating_sub(3600_000_000_000),
                10000,
                barrier_result,
                1,
                return_bps,
                0.8,
                50,
            );

            labels.push(label);
        }

        let weighted_samples = calculator.calculate_weights(&labels)?;

        assert_eq!(weighted_samples.len(), labels.len());

        // All weights should be positive
        for sample in &weighted_samples {
            assert!(sample.weight > 0.0);
            assert!(!sample.features.is_empty());
            assert_eq!(sample.features.len(), 3);
        }

        // Samples should have the expected structure
        assert_eq!(weighted_samples.len(), labels.len());

        Ok(())
    }
}