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

//! Sample Weights Calculator
//!
//! Implements sample weighting for addressing label imbalance and temporal decay,
//! based on MLFinLab methodology to reduce overfitting.
//!
//! ## Weighting Schemes
//!
//! 1. **Temporal Decay**: Recent samples weighted higher using exponential decay
//!    - Weight(t) = decay_factor^(days_old)
//!    - Typical decay_factor: 0.95 per day
//!
//! 2. **Label Balancing**: Balance class distribution
//!    - Weight(label) = 1 / count(label)
//!    - Prevents model from favoring majority class
//!
//! 3. **Combined**: Both temporal decay and label balancing
//!    - Weight = temporal_weight * balance_weight
//!
//! ## Usage
//!
//! ```rust
//! use ml::features::sample_weights::{SampleWeightCalculator, WeightingScheme};
//! use ml::labeling::meta_labeling::primary_model::Label;
//! use chrono::Utc;
//!
//! let calculator = SampleWeightCalculator::new(
//!     0.95,                         // decay_factor
//!     WeightingScheme::Combined,    // scheme
//! );
//!
//! let labels = vec![Label::Buy, Label::Sell, Label::Hold];
//! let timestamps = vec![Utc::now(); 3];
//!
//! let weights = calculator.calculate(&labels, &timestamps)?;
//! // weights sum to 1.0, ready for model training
//! ```

use crate::labeling::meta_labeling::primary_model::Label;
use crate::MLError;
use chrono::{DateTime, Utc};
use std::collections::HashMap;

/// Weighting scheme for sample weight calculation
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WeightingScheme {
    /// Only temporal decay (recent samples weighted higher)
    TemporalDecay,
    /// Only label balancing (balance class distribution)
    LabelBalancing,
    /// Both temporal decay and label balancing
    Combined,
}

/// Sample weight calculator for ML training
///
/// Computes sample weights to address:
/// - Label imbalance (buy/sell/hold distribution)
/// - Temporal decay (recent samples more relevant)
/// - Numerical stability (normalized weights sum to 1.0)
#[derive(Debug, Clone)]
pub struct SampleWeightCalculator {
    /// Exponential decay factor per day (typically 0.95)
    /// - decay_factor < 1.0: past weighted less (typical)
    /// - decay_factor = 1.0: no temporal weighting
    /// - decay_factor > 1.0: future weighted less (unusual but valid)
    decay_factor: f64,

    /// Weighting scheme to apply
    scheme: WeightingScheme,
}

impl SampleWeightCalculator {
    /// Create a new sample weight calculator
    ///
    /// # Arguments
    ///
    /// * `decay_factor` - Exponential decay per day (typically 0.95)
    /// * `scheme` - Weighting scheme to apply
    ///
    /// # Example
    ///
    /// ```rust
    /// use ml::features::sample_weights::{SampleWeightCalculator, WeightingScheme};
    ///
    /// let calculator = SampleWeightCalculator::new(0.95, WeightingScheme::Combined);
    /// ```
    pub fn new(decay_factor: f64, scheme: WeightingScheme) -> Self {
        Self {
            decay_factor,
            scheme,
        }
    }

    /// Calculate sample weights
    ///
    /// # Arguments
    ///
    /// * `labels` - Label for each sample (Buy/Sell/Hold)
    /// * `timestamps` - Timestamp for each sample
    ///
    /// # Returns
    ///
    /// Vector of weights normalized to sum to 1.0, or error if inputs are invalid
    ///
    /// # Errors
    ///
    /// - Empty inputs
    /// - Mismatched lengths
    /// - Invalid decay factor (0.0 or negative)
    ///
    /// # Example
    ///
    /// ```rust
    /// use ml::features::sample_weights::{SampleWeightCalculator, WeightingScheme};
    /// use ml::labeling::meta_labeling::primary_model::Label;
    /// use chrono::Utc;
    ///
    /// let calculator = SampleWeightCalculator::new(0.95, WeightingScheme::Combined);
    /// let labels = vec![Label::Buy, Label::Sell, Label::Hold];
    /// let timestamps = vec![Utc::now(); 3];
    ///
    /// let weights = calculator.calculate(&labels, &timestamps)?;
    /// assert!((weights.iter().sum::<f64>() - 1.0).abs() < 1e-6);
    /// # Ok::<(), ml::MLError>(())
    /// ```
    pub fn calculate(
        &self,
        labels: &[Label],
        timestamps: &[DateTime<Utc>],
    ) -> Result<Vec<f64>, MLError> {
        // Validate inputs
        if labels.is_empty() || timestamps.is_empty() {
            return Err(MLError::ConfigError("Labels and timestamps cannot be empty".to_owned()));
        }

        if labels.len() != timestamps.len() {
            return Err(MLError::ConfigError(format!(
                    "Labels length ({}) must match timestamps length ({})",
                    labels.len(),
                    timestamps.len()
                )));
        }

        if self.decay_factor <= 0.0 {
            return Err(MLError::ConfigError(format!("Decay factor must be positive, got {}", self.decay_factor)));
        }

        let n = labels.len();

        // Initialize weights to 1.0
        let mut weights = vec![1.0; n];

        // Apply temporal decay if needed
        if matches!(
            self.scheme,
            WeightingScheme::TemporalDecay | WeightingScheme::Combined
        ) {
            self.apply_temporal_decay(&mut weights, timestamps)?;
        }

        // Apply label balancing if needed
        if matches!(
            self.scheme,
            WeightingScheme::LabelBalancing | WeightingScheme::Combined
        ) {
            self.apply_label_balancing(&mut weights, labels)?;
        }

        // Normalize weights to sum to 1.0
        self.normalize_weights(&mut weights)?;

        Ok(weights)
    }

    /// Apply temporal decay to weights
    fn apply_temporal_decay(
        &self,
        weights: &mut [f64],
        timestamps: &[DateTime<Utc>],
    ) -> Result<(), MLError> {
        // Find the latest timestamp (most recent)
        let latest_time = timestamps
            .iter()
            .max()
            .ok_or_else(|| MLError::ConfigError("No timestamps provided".to_owned()))?;

        // Apply exponential decay based on age in days
        for (weight, timestamp) in weights.iter_mut().zip(timestamps.iter()) {
            let duration = *latest_time - *timestamp;
            let days_old = duration.num_days() as f64;

            // Weight = decay_factor^days_old
            // For decay_factor = 0.95, this means:
            // - 1 day old: weight = 0.95
            // - 2 days old: weight = 0.95^2 = 0.9025
            // - 30 days old: weight = 0.95^30 ≈ 0.215
            let decay_weight = self.decay_factor.powf(days_old);

            *weight *= decay_weight;
        }

        Ok(())
    }

    /// Apply label balancing to weights
    fn apply_label_balancing(&self, weights: &mut [f64], labels: &[Label]) -> Result<(), MLError> {
        // Count occurrences of each label
        let mut label_counts: HashMap<Label, usize> = HashMap::new();
        for label in labels {
            *label_counts.entry(*label).or_insert(0) += 1;
        }

        // Apply inverse frequency weighting
        // Weight = 1 / count(label)
        // This ensures that:
        // - Rare labels get higher weight
        // - Common labels get lower weight
        // - Total weight per class is approximately equal
        for (weight, label) in weights.iter_mut().zip(labels.iter()) {
            let count = label_counts
                .get(label)
                .ok_or_else(|| MLError::ConfigError(format!("Label {:?} not found in counts", label)))?;

            let balance_factor = 1.0 / (*count as f64);
            *weight *= balance_factor;
        }

        Ok(())
    }

    /// Normalize weights to sum to 1.0
    fn normalize_weights(&self, weights: &mut [f64]) -> Result<(), MLError> {
        let sum: f64 = weights.iter().sum();

        if sum <= 0.0 {
            return Err(MLError::ConfigError(format!("Weight sum must be positive, got {}", sum)));
        }

        // Normalize: divide each weight by the sum
        for weight in weights.iter_mut() {
            *weight /= sum;
        }

        Ok(())
    }

    /// Get the decay factor
    pub fn decay_factor(&self) -> f64 {
        self.decay_factor
    }

    /// Get the weighting scheme
    pub fn scheme(&self) -> WeightingScheme {
        self.scheme
    }
}

impl Default for SampleWeightCalculator {
    fn default() -> Self {
        Self {
            decay_factor: 0.95,
            scheme: WeightingScheme::Combined,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::Duration;

    #[test]
    fn test_basic_creation() {
        let calculator = SampleWeightCalculator::new(0.95, WeightingScheme::Combined);
        assert_eq!(calculator.decay_factor(), 0.95);
        assert_eq!(calculator.scheme(), WeightingScheme::Combined);
    }

    #[test]
    fn test_default() {
        let calculator = SampleWeightCalculator::default();
        assert_eq!(calculator.decay_factor(), 0.95);
        assert_eq!(calculator.scheme(), WeightingScheme::Combined);
    }

    #[test]
    fn test_single_sample() {
        let calculator = SampleWeightCalculator::new(0.95, WeightingScheme::Combined);
        let labels = vec![Label::Buy];
        let timestamps = vec![Utc::now()];

        let weights = calculator.calculate(&labels, &timestamps).unwrap();

        assert_eq!(weights.len(), 1);
        assert!((weights[0] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_temporal_decay_monotonic() {
        let calculator = SampleWeightCalculator::new(0.95, WeightingScheme::TemporalDecay);

        // Create samples with increasing age
        let labels = vec![Label::Buy; 5];
        let base_time = Utc::now();
        let timestamps: Vec<DateTime<Utc>> =
            (0..5).map(|i| base_time - Duration::days(i)).collect();

        let weights = calculator.calculate(&labels, &timestamps).unwrap();

        // Weights should be monotonically decreasing as samples get older
        for i in 0..weights.len() - 1 {
            assert!(
                weights[i] >= weights[i + 1],
                "Weight[{}] = {} should be >= Weight[{}] = {}",
                i,
                weights[i],
                i + 1,
                weights[i + 1]
            );
        }
    }

    #[test]
    fn test_label_balancing_effect() {
        let calculator = SampleWeightCalculator::new(1.0, WeightingScheme::LabelBalancing);

        // 3 Buy, 1 Sell, 1 Hold
        let labels = vec![Label::Buy, Label::Buy, Label::Buy, Label::Sell, Label::Hold];
        let timestamps = vec![Utc::now(); 5];

        let weights = calculator.calculate(&labels, &timestamps).unwrap();

        // Sell and Hold should have higher weights than Buy
        assert!(weights[3] > weights[0]); // Sell > Buy
        assert!(weights[4] > weights[0]); // Hold > Buy
    }

    #[test]
    fn test_normalization() {
        let schemes = vec![
            WeightingScheme::TemporalDecay,
            WeightingScheme::LabelBalancing,
            WeightingScheme::Combined,
        ];

        for scheme in schemes {
            let calculator = SampleWeightCalculator::new(0.95, scheme);
            let labels = vec![Label::Buy, Label::Sell, Label::Hold];
            let timestamps = vec![Utc::now(); 3];

            let weights = calculator.calculate(&labels, &timestamps).unwrap();

            let sum: f64 = weights.iter().sum();
            assert!(
                (sum - 1.0).abs() < 1e-10,
                "Weights should sum to 1.0 for scheme {:?}, got {}",
                scheme,
                sum
            );
        }
    }
}