foxhunt/crates/ml-dqn/src/evaluation/engine.rs

//! Evaluation Engine for DQN Backtesting
//!
//! Tracks positions, executes trades based on DQN actions, and records trade history.

use super::metrics::OHLCVBarF32;
use serde::{Deserialize, Serialize};

/// Trading action from DQN model
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Action {
    Buy = 0,
    Hold = 1,
    Sell = 2,
}

impl From<usize> for Action {
    fn from(action: usize) -> Self {
        match action {
            0 => Action::Buy,
            1 => Action::Hold,
            2 => Action::Sell,
            _ => Action::Hold, // Default to hold for invalid actions
        }
    }
}

/// Open position
#[derive(Debug, Clone)]
pub struct Position {
    pub entry_bar_idx: usize,
    pub entry_price: f32,
    pub direction: PositionDirection,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PositionDirection {
    Long,
    Short,
}

/// Completed trade
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Trade {
    pub entry_bar_idx: usize,
    pub exit_bar_idx: usize,
    pub entry_price: f32,
    pub exit_price: f32,
    pub direction: String,
    pub pnl: f32,
}

/// Evaluation engine that processes DQN actions and tracks positions
#[derive(Debug)]
pub struct EvaluationEngine {
    pub current_position: Option<Position>,
    pub trades: Vec<Trade>,
    pub initial_capital: f32,
    pub action_counts: [usize; 3], // [buy, hold, sell]
    /// Kelly fraction for position sizing (1.0 = full size, 0.5 = half size)
    pub kelly_fraction: f64,
    /// Current position exposure for factored evaluation (-1.0 to +1.0)
    pub current_exposure: f64,
    /// Entry price for current exposure-based position
    pub exposure_entry_price: f32,
    /// Bar index where current exposure was first entered
    pub exposure_entry_bar: usize,
    /// Override fee rate for all trades (None = use action's `transaction_cost()`)
    fee_rate_override: Option<f64>,
    /// Running equity (`initial_capital` + sum of all trade `PnLs`)
    pub running_equity: f32,
    /// True when running equity hit zero — no further trades are executed
    pub margin_called: bool,
}

impl EvaluationEngine {
    /// Create new evaluation engine
    /// 
    /// # Arguments
    /// * `initial_capital` - Starting capital for backtest
    /// * `kelly_fraction` - Position sizing multiplier (default: 1.0 = full size)
    pub const fn new_with_kelly(initial_capital: f32, kelly_fraction: f64) -> Self {
        Self {
            current_position: None,
            trades: Vec::new(),
            initial_capital,
            action_counts: [0, 0, 0],
            kelly_fraction,
            current_exposure: 0.0,
            exposure_entry_price: 0.0,
            exposure_entry_bar: 0,
            fee_rate_override: None,
            running_equity: initial_capital,
            margin_called: false,
        }
    }

    /// Create evaluation engine with explicit fee rate override (in decimal, e.g. 0.00001 for 0.1 bps)
    pub const fn new_with_fee_rate(initial_capital: f32, kelly_fraction: f64, fee_rate: f64) -> Self {
        Self {
            current_position: None,
            trades: Vec::new(),
            initial_capital,
            action_counts: [0, 0, 0],
            kelly_fraction,
            current_exposure: 0.0,
            exposure_entry_price: 0.0,
            exposure_entry_bar: 0,
            fee_rate_override: Some(fee_rate),
            running_equity: initial_capital,
            margin_called: false,
        }
    }

    /// Create new evaluation engine with default Kelly fraction (1.0 = full size)
    pub const fn new(initial_capital: f32) -> Self {
        Self::new_with_kelly(initial_capital, 1.0)
    }

    /// Process a single bar with DQN action
    ///
    /// # Arguments
    /// * `bar_idx` - Index of current bar in the dataset
    /// * `bar` - Current OHLCV bar
    /// * `action` - Action selected by DQN model
    pub fn process_bar(&mut self, bar_idx: usize, bar: &OHLCVBarF32, action: Action) {
        if self.margin_called {
            return;
        }
        // Update action counts
        self.action_counts[action as usize] += 1;

        match action {
            Action::Buy => {
                // If no position or short position, open long
                if let Some(pos) = &self.current_position {
                    if pos.direction == PositionDirection::Short {
                        // Close short position
                        self.close_position(bar_idx, bar);
                    }
                }

                // Open new long position
                if self.current_position.is_none() {
                    self.current_position = Some(Position {
                        entry_bar_idx: bar_idx,
                        entry_price: bar.close,
                        direction: PositionDirection::Long,
                    });
                }
            },

            Action::Sell => {
                // If no position or long position, open short
                if let Some(pos) = &self.current_position {
                    if pos.direction == PositionDirection::Long {
                        // Close long position
                        self.close_position(bar_idx, bar);
                    }
                }

                // Open new short position
                if self.current_position.is_none() {
                    self.current_position = Some(Position {
                        entry_bar_idx: bar_idx,
                        entry_price: bar.close,
                        direction: PositionDirection::Short,
                    });
                }
            },

            Action::Hold => {
                // Do nothing, maintain current position
            },
        }
    }

    /// Close current position and record trade
    pub fn close_position(&mut self, exit_bar_idx: usize, exit_bar: &OHLCVBarF32) {
        if let Some(pos) = self.current_position.take() {
            // Calculate base PnL (for 1 contract)
            let base_pnl = match pos.direction {
                PositionDirection::Long => {
                    // Long: profit when price goes up
                    exit_bar.close - pos.entry_price
                },
                PositionDirection::Short => {
                    // Short: profit when price goes down
                    pos.entry_price - exit_bar.close
                },
            };

            // Apply Kelly scaling to gross PnL (Kelly fraction scales position size)
            // Example: Kelly=0.5 means half position, so gross PnL is also halved
            let kelly_scaled_gross_pnl = base_pnl * self.kelly_fraction as f32;

            // Calculate transaction costs (Bug #3 fix: Net P&L calculation)
            // Assumes market orders (0.15% fee) on both entry and exit
            const MARKET_ORDER_FEE: f32 = 0.0015; // 0.15%

            // Transaction costs scale with Kelly fraction (smaller position = smaller costs)
            let entry_cost = pos.entry_price * self.kelly_fraction as f32 * MARKET_ORDER_FEE;
            let exit_cost = exit_bar.close * self.kelly_fraction as f32 * MARKET_ORDER_FEE;
            let total_transaction_cost = entry_cost + exit_cost;

            // Net P&L = Gross P&L - Transaction Costs
            let net_pnl = kelly_scaled_gross_pnl - total_transaction_cost;

            let trade = Trade {
                entry_bar_idx: pos.entry_bar_idx,
                exit_bar_idx,
                entry_price: pos.entry_price,
                exit_price: exit_bar.close,
                direction: match pos.direction {
                    PositionDirection::Long => "long".to_owned(),
                    PositionDirection::Short => "short".to_owned(),
                },
                pnl: net_pnl, // Store NET P&L (after costs)
            };

            self.trades.push(trade);
            self.running_equity += net_pnl;

            if self.running_equity <= 0.0 {
                self.margin_called = true;
            }
        }
    }

    /// Process a bar using the full 45-action factored space.
    ///
    /// Tracks continuous exposure (-1.0 to +1.0) and generates trades on
    /// any exposure change (including partial: Long100->Long50 = sell 0.5).
    /// Transaction costs use the `FactoredAction`'s order type.
    pub fn process_bar_factored(
        &mut self,
        bar_idx: usize,
        bar: &OHLCVBarF32,
        action: &crate::action_space::FactoredAction,
    ) {
        // Margin call: equity depleted — force-flat and stop trading
        if self.margin_called {
            return;
        }

        let target = action.target_exposure(); // -1.0, -0.5, 0.0, +0.5, +1.0
        let delta = target - self.current_exposure;

        // Update legacy action counts for compatibility with metrics
        if target > 0.0 {
            self.action_counts[0] += 1; // buy
        } else if target < 0.0 {
            self.action_counts[2] += 1; // sell
        } else {
            self.action_counts[1] += 1; // hold
        }

        const EPSILON: f64 = 1e-6;
        if delta.abs() < EPSILON {
            return; // No position change
        }

        // Record trade for the exposure change
        let fee_rate = self.fee_rate_override.unwrap_or_else(|| action.transaction_cost());

        // PnL from the portion being closed (if reducing or reversing)
        let closing_size = if self.current_exposure.abs() > EPSILON
            && delta.signum() != self.current_exposure.signum()
        {
            // Closing part (or all) of existing position
            self.current_exposure.abs().min(delta.abs())
        } else {
            0.0
        };

        if closing_size > EPSILON {
            let price_diff = bar.close - self.exposure_entry_price;
            let direction_sign = if self.current_exposure > 0.0 {
                1.0_f32
            } else {
                -1.0_f32
            };
            let gross_pnl =
                price_diff * direction_sign * (closing_size * self.kelly_fraction) as f32;
            let tx_cost = (self.exposure_entry_price.abs() + bar.close.abs()) as f64
                * 0.5
                * closing_size
                * self.kelly_fraction
                * fee_rate;

            let net_pnl = gross_pnl - tx_cost as f32;
            self.trades.push(Trade {
                entry_bar_idx: self.exposure_entry_bar,
                exit_bar_idx: bar_idx,
                entry_price: self.exposure_entry_price,
                exit_price: bar.close,
                direction: if self.current_exposure > 0.0 {
                    "long".to_owned()
                } else {
                    "short".to_owned()
                },
                pnl: net_pnl,
            });
            self.running_equity += net_pnl;

            // Margin call: equity depleted — record but stop further trading
            if self.running_equity <= 0.0 {
                self.margin_called = true;
                self.current_exposure = 0.0;
                return;
            }
        }

        // Update exposure state
        if target.abs() > EPSILON {
            // Opening or adjusting -- reset entry if crossing zero or first entry
            if self.current_exposure.abs() < EPSILON
                || target.signum() != self.current_exposure.signum()
            {
                self.exposure_entry_price = bar.close;
                self.exposure_entry_bar = bar_idx;
            }
        }
        self.current_exposure = target;
    }

    /// Close any remaining factored exposure at end of backtest
    pub fn close_factored_position(&mut self, bar_idx: usize, bar: &OHLCVBarF32) {
        const EPSILON: f64 = 1e-6;
        if self.current_exposure.abs() < EPSILON || self.margin_called {
            return;
        }
        let price_diff = bar.close - self.exposure_entry_price;
        let direction_sign = if self.current_exposure > 0.0 {
            1.0_f32
        } else {
            -1.0_f32
        };
        let size = self.current_exposure.abs() * self.kelly_fraction;
        let gross_pnl = price_diff * direction_sign * size as f32;
        let close_fee = self.fee_rate_override.unwrap_or(0.0015);
        let tx_cost = bar.close.abs() as f64 * size * close_fee;

        let net_pnl = gross_pnl - tx_cost as f32;
        self.trades.push(Trade {
            entry_bar_idx: self.exposure_entry_bar,
            exit_bar_idx: bar_idx,
            entry_price: self.exposure_entry_price,
            exit_price: bar.close,
            direction: if self.current_exposure > 0.0 {
                "long".to_owned()
            } else {
                "short".to_owned()
            },
            pnl: net_pnl,
        });
        self.running_equity += net_pnl;
        self.current_exposure = 0.0;

        if self.running_equity <= 0.0 {
            self.margin_called = true;
        }
    }

    /// Get action distribution summary
    pub fn get_action_distribution(&self) -> ActionDistribution {
        let total = self.action_counts.iter().sum::<usize>();
        let total_f64 = total as f64;

        ActionDistribution {
            buy_count: self.action_counts[0],
            hold_count: self.action_counts[1],
            sell_count: self.action_counts[2],
            buy_pct: if total > 0 {
                (self.action_counts[0] as f64 / total_f64) * 100.0
            } else {
                0.0
            },
            hold_pct: if total > 0 {
                (self.action_counts[1] as f64 / total_f64) * 100.0
            } else {
                0.0
            },
            sell_pct: if total > 0 {
                (self.action_counts[2] as f64 / total_f64) * 100.0
            } else {
                0.0
            },
        }
    }
}

/// Action distribution statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ActionDistribution {
    pub buy_count: usize,
    pub hold_count: usize,
    pub sell_count: usize,
    pub buy_pct: f64,
    pub hold_pct: f64,
    pub sell_pct: f64,
}

#[cfg(test)]
#[allow(clippy::unnecessary_map_or)]
mod tests {
    use super::*;
    use crate::action_space::{ExposureLevel, FactoredAction, OrderType, Urgency};

    fn bar(close: f32) -> OHLCVBarF32 {
        OHLCVBarF32 {
            timestamp: 0,
            open: close,
            high: close,
            low: close,
            close,
            volume: 0.0,
        }
    }

    fn market_action(exposure: ExposureLevel) -> FactoredAction {
        FactoredAction::new(exposure, OrderType::Market, Urgency::Normal)
    }

    #[test]
    fn factored_same_exposure_no_trade() {
        let mut engine = EvaluationEngine::new(10000.0);
        let b = bar(100.0);
        engine.process_bar_factored(0, &b, &market_action(ExposureLevel::Long100));
        engine.process_bar_factored(1, &b, &market_action(ExposureLevel::Long100));
        assert_eq!(
            engine.trades.len(),
            0,
            "Same exposure should generate no trades"
        );
    }

    #[test]
    fn factored_partial_close_generates_trade() {
        let mut engine = EvaluationEngine::new(10000.0);
        engine.process_bar_factored(0, &bar(100.0), &market_action(ExposureLevel::Long100));
        engine.process_bar_factored(1, &bar(110.0), &market_action(ExposureLevel::Long50));
        assert_eq!(
            engine.trades.len(),
            1,
            "Long100->Long50 should generate 1 trade"
        );
        assert!(
            engine.trades.first().map_or(false, |t| t.pnl > 0.0),
            "Price went up on long = profit"
        );
        assert!((engine.current_exposure - 0.5).abs() < 1e-6);
    }

    #[test]
    fn factored_reversal_generates_trade() {
        let mut engine = EvaluationEngine::new(10000.0);
        engine.process_bar_factored(0, &bar(100.0), &market_action(ExposureLevel::Long100));
        engine.process_bar_factored(1, &bar(105.0), &market_action(ExposureLevel::Short100));
        assert!(
            !engine.trades.is_empty(),
            "Reversal should generate at least 1 trade"
        );
        assert!((engine.current_exposure - (-1.0)).abs() < 1e-6);
    }

    #[test]
    fn factored_flat_from_long_closes() {
        let mut engine = EvaluationEngine::new(10000.0);
        engine.process_bar_factored(0, &bar(100.0), &market_action(ExposureLevel::Long100));
        engine.process_bar_factored(1, &bar(95.0), &market_action(ExposureLevel::Flat));
        assert_eq!(engine.trades.len(), 1);
        assert!(
            engine.trades.first().map_or(false, |t| t.pnl < 0.0),
            "Price went down on long = loss"
        );
        assert!(engine.current_exposure.abs() < 1e-6);
    }

    #[test]
    fn factored_close_at_end() {
        let mut engine = EvaluationEngine::new(10000.0);
        engine.process_bar_factored(0, &bar(100.0), &market_action(ExposureLevel::Short50));
        engine.close_factored_position(1, &bar(90.0));
        assert_eq!(engine.trades.len(), 1);
        assert!(
            engine.trades.first().map_or(false, |t| t.pnl > 0.0),
            "Price down on short = profit"
        );
        assert!(engine.current_exposure.abs() < 1e-6);
    }

    #[test]
    fn factored_all_buy_still_one_trade_at_close() {
        let mut engine = EvaluationEngine::new(10000.0);
        for i in 0..100 {
            engine.process_bar_factored(
                i,
                &bar(100.0 + i as f32),
                &market_action(ExposureLevel::Long100),
            );
        }
        assert_eq!(engine.trades.len(), 0);
        engine.close_factored_position(100, &bar(200.0));
        assert_eq!(engine.trades.len(), 1);
    }

    #[test]
    fn factored_alternating_generates_many_trades() {
        let mut engine = EvaluationEngine::new(10000.0);
        for i in 0..10 {
            let action = if i % 2 == 0 {
                market_action(ExposureLevel::Long100)
            } else {
                market_action(ExposureLevel::Short100)
            };
            engine.process_bar_factored(i, &bar(100.0), &action);
        }
        assert!(
            engine.trades.len() >= 9,
            "Alternating should generate many trades: got {}",
            engine.trades.len()
        );
    }
}