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

//! Position-Aware Feature Extraction for RL Agents
//!
//! Extracts 3 features that give the RL agent awareness of its own position:
//! - `unrealized_pnl`: normalized by EMA volatility
//! - `bars_in_position`: log-scaled to [0, 1]
//! - `cost_basis`: relative to current price in bps, clamped to [-500, 500]
//!
//! These features are appended to the state vector after the 40 market features
//! (producing a 43-dim state with portfolio, or 51-dim with OFI).

/// Maximum number of bars for log scaling (1 trading day at 1-min resolution).
const MAX_BARS_LOG: f64 = 390.0;

/// Cost basis clamp bound in basis points.
const COST_BASIS_CLAMP_BPS: f64 = 500.0;

/// Default EMA span for volatility estimation.
const DEFAULT_VOL_SPAN: usize = 100;

/// Minimum volatility floor to prevent division by zero.
const MIN_VOLATILITY: f64 = 1e-10;

/// Extracts position-aware features for RL agents.
///
/// Maintains an EMA of absolute returns to normalize unrealized `PnL`,
/// ensuring the agent perceives `PnL` relative to recent market movement.
#[derive(Debug, Clone)]
pub struct PositionFeatures {
    /// EMA smoothing factor: alpha = 2 / (span + 1)
    alpha: f64,
    /// Current EMA of absolute returns (volatility proxy)
    ema_volatility: Option<f64>,
}

impl Default for PositionFeatures {
    fn default() -> Self {
        Self::new(DEFAULT_VOL_SPAN)
    }
}

impl PositionFeatures {
    /// Create a new `PositionFeatures` extractor with the given EMA span.
    ///
    /// # Arguments
    ///
    /// * `span` - Number of periods for the EMA volatility estimate.
    ///   A span of 100 gives alpha = 2/101 ~ 0.0198.
    #[must_use]
    pub fn new(span: usize) -> Self {
        let alpha = 2.0 / (span as f64 + 1.0);
        Self {
            alpha,
            ema_volatility: None,
        }
    }

    /// Update the EMA volatility estimate with a new absolute return observation.
    ///
    /// Call this once per bar with `|close_t / close_{t-1} - 1|`.
    pub fn update_volatility(&mut self, abs_return: f64) {
        let val = abs_return.abs(); // defensive: ensure positive
        self.ema_volatility = Some(match self.ema_volatility {
            Some(prev) => self.alpha * val + (1.0 - self.alpha) * prev,
            None => val, // seed with first observation
        });
    }

    /// Extract 3 position-aware features.
    ///
    /// Returns `[unrealized_pnl_norm, bars_in_position_norm, cost_basis_bps_norm]`.
    ///
    /// # Arguments
    ///
    /// * `position_size` - Current position size (positive = long, negative = short, 0 = flat).
    /// * `unrealized_pnl` - Dollar `PnL` of the open position.
    /// * `bars_held` - Number of bars the current position has been held.
    /// * `entry_price` - Average entry price of the current position.
    /// * `current_price` - Current mid/close price.
    ///
    /// # Feature definitions
    ///
    /// 1. **`unrealized_pnl_norm`**: `unrealized_pnl / ema_volatility` (0 when flat or no vol).
    /// 2. **`bars_in_position_norm`**: `ln(1 + bars_held) / ln(1 + MAX_BARS)`, clamped to [0, 1].
    /// 3. **`cost_basis_bps_norm`**: `(current_price / entry_price - 1) * 10_000`, clamped to [-500, 500],
    ///    then scaled to [-1, 1] by dividing by 500. Sign is flipped for short positions.
    #[must_use]
    pub fn extract(
        &self,
        position_size: f64,
        unrealized_pnl: f64,
        bars_held: u64,
        entry_price: f64,
        current_price: f64,
    ) -> Vec<f64> {
        let is_flat = position_size.abs() < f64::EPSILON;

        // Feature 1: unrealized PnL normalized by EMA volatility
        let pnl_norm = if is_flat {
            0.0
        } else {
            let vol = self.ema_volatility.unwrap_or(MIN_VOLATILITY).max(MIN_VOLATILITY);
            unrealized_pnl / vol
        };

        // Feature 2: bars in position, log-scaled to [0, 1]
        let bars_norm = if is_flat {
            0.0
        } else {
            let log_bars = (1.0 + bars_held as f64).ln();
            let log_max = (1.0 + MAX_BARS_LOG).ln();
            (log_bars / log_max).clamp(0.0, 1.0)
        };

        // Feature 3: cost basis relative to current price in bps, scaled to [-1, 1]
        let cost_basis_norm = if is_flat || entry_price.abs() < f64::EPSILON {
            0.0
        } else {
            let raw_bps = (current_price / entry_price - 1.0) * 10_000.0;
            // Flip sign for short positions: a price increase is bad for shorts
            let directional_bps = if position_size < 0.0 { -raw_bps } else { raw_bps };
            directional_bps.clamp(-COST_BASIS_CLAMP_BPS, COST_BASIS_CLAMP_BPS) / COST_BASIS_CLAMP_BPS
        };

        vec![pnl_norm, bars_norm, cost_basis_norm]
    }
}

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

    const EPSILON: f64 = 1e-9;

    #[test]
    fn test_flat_position_returns_zeros() {
        let pf = PositionFeatures::default();
        let features = pf.extract(0.0, 0.0, 0, 0.0, 100.0);
        assert_eq!(features.len(), 3);
        assert!((features[0]).abs() < EPSILON, "pnl should be 0 when flat");
        assert!((features[1]).abs() < EPSILON, "bars should be 0 when flat");
        assert!((features[2]).abs() < EPSILON, "cost basis should be 0 when flat");
    }

    #[test]
    fn test_long_position_positive_pnl() {
        let mut pf = PositionFeatures::new(100);
        // Seed volatility with a few updates
        for _ in 0..10 {
            pf.update_volatility(0.01); // 1% absolute return
        }

        let features = pf.extract(
            1.0,    // long 1 unit
            500.0,  // $500 unrealized profit
            60,     // held for 60 bars
            100.0,  // entered at 100
            105.0,  // now at 105
        );

        assert_eq!(features.len(), 3);
        // PnL normalized by vol should be positive
        assert!(features[0] > 0.0, "long + profit => positive pnl_norm");
        // 60 bars => ln(61)/ln(391) ~ 4.11/5.97 ~ 0.689
        assert!(features[1] > 0.5 && features[1] < 0.8, "60 bars should be ~0.69, got {}", features[1]);
        // (105/100 - 1) * 10000 = 500 bps => 500/500 = 1.0
        assert!((features[2] - 1.0).abs() < EPSILON, "cost basis should be 1.0, got {}", features[2]);
    }

    #[test]
    fn test_short_position_flips_cost_basis() {
        let mut pf = PositionFeatures::new(100);
        pf.update_volatility(0.01);

        let features = pf.extract(
            -1.0,   // short 1 unit
            200.0,  // $200 unrealized profit (price went down)
            30,     // held for 30 bars
            100.0,  // entered at 100
            98.0,   // now at 98 (good for short)
        );

        assert_eq!(features.len(), 3);
        // PnL normalized by vol should be positive (profitable short)
        assert!(features[0] > 0.0, "short + profit => positive pnl_norm");
        // (98/100 - 1) * 10000 = -200 bps, then flipped for short => +200 bps => 200/500 = 0.4
        assert!(
            (features[2] - 0.4).abs() < 0.01,
            "short cost basis should be ~0.4, got {}",
            features[2]
        );
    }

    #[test]
    fn test_log_scaling_bars() {
        let pf = PositionFeatures::default();

        // 1 bar: ln(2)/ln(391) ~ 0.693/5.97 ~ 0.116
        let f1 = pf.extract(1.0, 0.0, 1, 100.0, 100.0);
        assert!(f1[1] > 0.1 && f1[1] < 0.15, "1 bar ~ 0.116, got {}", f1[1]);

        // 390 bars (full day): ln(391)/ln(391) = 1.0
        let f390 = pf.extract(1.0, 0.0, 390, 100.0, 100.0);
        assert!((f390[1] - 1.0).abs() < EPSILON, "390 bars should be 1.0, got {}", f390[1]);

        // 1000 bars (> max): should clamp to 1.0
        let f1000 = pf.extract(1.0, 0.0, 1000, 100.0, 100.0);
        assert!((f1000[1] - 1.0).abs() < EPSILON, "1000 bars should clamp to 1.0, got {}", f1000[1]);
    }

    #[test]
    fn test_cost_basis_clamping() {
        let pf = PositionFeatures::default();

        // Extreme price move: 100 -> 200 => (200/100 - 1) * 10000 = 10000 bps
        // Should clamp to 500 bps => 500/500 = 1.0
        let f_up = pf.extract(1.0, 0.0, 10, 100.0, 200.0);
        assert!(
            (f_up[2] - 1.0).abs() < EPSILON,
            "extreme up should clamp to 1.0, got {}",
            f_up[2]
        );

        // Extreme price drop: 100 -> 50 => (50/100 - 1) * 10000 = -5000 bps
        // Should clamp to -500 bps => -500/500 = -1.0
        let f_down = pf.extract(1.0, 0.0, 10, 100.0, 50.0);
        assert!(
            (f_down[2] - (-1.0)).abs() < EPSILON,
            "extreme down should clamp to -1.0, got {}",
            f_down[2]
        );
    }

    #[test]
    fn test_ema_volatility_update() {
        let mut pf = PositionFeatures::new(100);
        assert!(pf.ema_volatility.is_none());

        // First update seeds the EMA
        pf.update_volatility(0.02);
        assert!((pf.ema_volatility.unwrap() - 0.02).abs() < EPSILON);

        // Second update blends: alpha * 0.01 + (1-alpha) * 0.02
        // alpha = 2/101 ~ 0.0198
        pf.update_volatility(0.01);
        let alpha = 2.0 / 101.0;
        let expected = alpha * 0.01 + (1.0 - alpha) * 0.02;
        assert!(
            (pf.ema_volatility.unwrap() - expected).abs() < EPSILON,
            "EMA should be {}, got {}",
            expected,
            pf.ema_volatility.unwrap()
        );
    }

    #[test]
    fn test_zero_entry_price_returns_zero_cost_basis() {
        let pf = PositionFeatures::default();
        let features = pf.extract(1.0, 100.0, 10, 0.0, 100.0);
        assert!((features[2]).abs() < EPSILON, "zero entry price => 0 cost basis");
    }

    #[test]
    fn test_no_volatility_uses_floor() {
        let pf = PositionFeatures::default(); // no update_volatility calls
        let features = pf.extract(1.0, 100.0, 10, 100.0, 101.0);
        // PnL should be 100 / 1e-10 = 1e12 (huge, but finite)
        assert!(features[0].is_finite(), "pnl_norm should be finite");
        assert!(features[0] > 0.0, "pnl_norm should be positive");
    }
}