//! Portfolio state tracking for PPO training
//!
//! This module provides portfolio state management for the PPO agent, tracking:
//! - Portfolio value (cash + unrealized P&L)
//! - Position size (signed: +Long, -Short, 0 for flat)
//! - Bid-ask spread (for transaction costs)
//!
//! The `PortfolioTracker` is used by PPO trainers to provide portfolio features
//! to the reward function, enabling P&L-based reward calculations.
//!
//! This implementation is adapted from DQN's `PortfolioTracker` to work with
//! PPO's action space (simple usize indices: 0=BUY, 1=SELL, 2=HOLD).

use tracing::warn;

/// Portfolio state tracker for PPO training
///
/// Tracks cash, position size, and spread to provide portfolio features
/// for reward calculations. The tracker maintains state across training
/// steps within an epoch, and can be reset at epoch boundaries.
#[derive(Debug, Clone)]
pub struct PortfolioTracker {
    /// Current cash balance
    cash: f32,
    /// Current position size (positive = long, negative = short, 0 = flat)
    position_size: f32,
    /// Entry price for current position
    position_entry_price: f32,
    /// Initial capital (for reset)
    initial_capital: f32,
    /// Average bid-ask spread (estimated from historical data)
    avg_spread: f32,
    /// Last observed price (for parameter-less `total_value()` calls)
    last_price: f32,
    /// Cash reserve requirement as percentage of portfolio value (0-100)
    cash_reserve_percent: f32,
    /// Cumulative transaction costs
    cumulative_transaction_costs: f32,
}

impl PortfolioTracker {
    /// Create a new portfolio tracker
    ///
    /// # Arguments
    ///
    /// * `initial_capital` - Starting cash balance (e.g., 10,000.0)
    /// * `avg_spread` - Average bid-ask spread as a fraction (e.g., 0.0001 = 1 basis point)
    /// * `cash_reserve_percent` - Cash reserve requirement as percentage of portfolio value (0-100, default: 0)
    ///
    /// # Example
    ///
    /// ```
    /// use ml::ppo::portfolio_tracker::PortfolioTracker;
    ///
    /// let tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);
    /// ```
    pub const fn new(initial_capital: f32, avg_spread: f32, cash_reserve_percent: f64) -> Self {
        Self {
            cash: initial_capital,
            position_size: 0.0,
            position_entry_price: 0.0,
            initial_capital,
            avg_spread,
            last_price: 0.0,
            cash_reserve_percent: cash_reserve_percent as f32,
            cumulative_transaction_costs: 0.0,
        }
    }

    /// Create a new portfolio tracker with default spread
    ///
    /// This is a convenience constructor.
    /// Uses a default spread of 0.0001 (1 basis point).
    ///
    /// # Arguments
    ///
    /// * `initial_capital` - Starting cash balance (e.g., 10,000.0)
    ///
    /// # Example
    ///
    /// ```
    /// use ml::ppo::portfolio_tracker::PortfolioTracker;
    ///
    /// let tracker = PortfolioTracker::with_default_spread(10_000.0);
    /// ```
    pub const fn with_default_spread(initial_capital: f32) -> Self {
        Self::new(initial_capital, 0.0001, 0.0)
    }

    /// Get portfolio features for PPO reward calculation
    ///
    /// Returns a 3-element array:
    /// - `[0]`: Portfolio value (cash + unrealized P&L)
    /// - `[1]`: Position size (signed: +Long, -Short, 0 for flat)
    /// - `[2]`: Bid-ask spread (for transaction costs)
    ///
    /// # Arguments
    ///
    /// * `current_price` - Current market price for calculating unrealized P&L
    ///
    /// # Example
    ///
    /// ```
    /// use ml::ppo::portfolio_tracker::PortfolioTracker;
    ///
    /// let tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);
    /// let features = tracker.get_portfolio_features(100.0);
    /// assert_eq!(features[0], 1.0); // Normalized portfolio value (1.0 = initial capital)
    /// assert_eq!(features[1], 0.0); // No position
    /// assert_eq!(features[2], 0.0001); // Spread
    /// ```
    pub fn get_portfolio_features(&self, current_price: f32) -> [f32; 3] {
        let portfolio_value = self.get_portfolio_value(current_price);

        // Normalize portfolio value by initial capital (1.0 = initial capital)
        // Example: 10,100 / 10,000 = 1.01 (1% gain)
        let normalized_value = portfolio_value / self.initial_capital;

        // Normalize position size: Assume max position = initial_capital / current_price
        // Example: For $10,000 capital at $100/contract: max = 100 contracts
        // Position of 10 contracts = 10/100 = 0.1 (10% exposure)
        let max_position = if current_price > 0.0 {
            self.initial_capital / current_price
        } else {
            1.0 // Fallback to avoid division by zero
        };
        let normalized_position = self.position_size / max_position;

        [
            normalized_value,    // [0] Normalized portfolio value (1.0 = initial capital)
            normalized_position, // [1] Normalized position size (1.0 = max exposure)
            self.avg_spread,     // [2] Spread (already small, doesn't need normalization)
        ]
    }

    /// Execute PPO action and update portfolio state
    ///
    /// # Arguments
    ///
    /// * `action_index` - PPO action index (0=BUY, 1=SELL, 2=HOLD)
    /// * `price` - Current market price
    /// * `position_units` - Number of units to trade (e.g., 10.0 contracts)
    ///
    /// # Action Behavior
    ///
    /// - **0 (BUY)**: Opens long position if flat, closes short position, or reverses from long to bigger long
    /// - **1 (SELL)**: Opens short position if flat, closes long position, or reverses from short to bigger short
    /// - **2 (HOLD)**: No change to portfolio state
    ///
    /// # Position Reversal
    ///
    /// Position reversal occurs when the action would flip the position sign (e.g., long to short).
    /// This is handled in two phases:
    /// - **Phase 1**: Close current position (realize P&L)
    /// - **Phase 2**: Open new position in opposite direction (may be partial if insufficient cash)
    ///
    /// # Example
    ///
    /// ```
    /// use ml::ppo::portfolio_tracker::PortfolioTracker;
    ///
    /// let mut tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);
    /// tracker.execute_ppo_action(0, 100.0, 10.0); // BUY 10 contracts
    /// assert_eq!(tracker.current_position(), 10.0);
    /// assert_eq!(tracker.cash_balance(), 9_000.0); // 10_000 - (10 * 100)
    /// ```
    pub fn execute_ppo_action(&mut self, action_index: usize, price: f32, position_units: f32) {
        // Map PPO action index to trading action
        // 0 = BUY, 1 = SELL, 2 = HOLD (3-action space)
        // This will be extended to 45 actions later
        match action_index {
            0 => {
                // BUY
                // Check for position reversal (long position being sold beyond flat)
                if self.position_size < 0.0 && position_units > self.position_size.abs() {
                    // Reversal detected: short -> long
                    self.handle_reversal(price, position_units, true);
                    return;
                }

                if self.position_size == 0.0 {
                    // P2-B: Cash Reserve Check (BUY trades only)
                    if self.cash_reserve_percent > 0.0 {
                        let trade_cost = position_units * price;
                        let portfolio_value = self.get_portfolio_value(price);
                        let reserve_required = portfolio_value * (self.cash_reserve_percent / 100.0);
                        let cash_after_trade = self.cash - trade_cost;

                        if cash_after_trade < reserve_required {
                            // Calculate affordable position (respecting reserve)
                            let affordable_cash = (self.cash - reserve_required).max(0.0);
                            let affordable_contracts = (affordable_cash / price).floor();

                            if affordable_contracts <= 0.0 {
                                warn!(
                                    "P2-B: BUY REJECTED - cash reserve violated. Cash: ${:.2}, Reserve: ${:.2} ({:.1}%), Trade cost: ${:.2}",
                                    self.cash, reserve_required, self.cash_reserve_percent, trade_cost
                                );
                                return;
                            }

                            warn!(
                                "P2-B: BUY REDUCED - cash reserve enforced. Requested: {:.1}, Reduced to: {:.1}",
                                position_units, affordable_contracts
                            );

                            // Execute with reduced position
                            self.position_size = affordable_contracts;
                            self.position_entry_price = price;
                            self.cash -= affordable_contracts * price;
                            return;
                        }
                    }

                    // Open long position (no cash reserve violation)
                    self.position_size = position_units;
                    self.position_entry_price = price;
                    self.cash -= position_units * price;
                } else if self.position_size < 0.0 {
                    // Close short position (buy to cover)
                    self.cash += self.position_size * price; // position_size is negative
                    self.position_size = 0.0;
                    self.position_entry_price = 0.0;
                } else {
                    // Already long, do nothing
                }
            }
            1 => {
                // SELL
                // Check for position reversal (short position being bought beyond flat)
                if self.position_size > 0.0 && position_units > self.position_size {
                    // Reversal detected: long -> short
                    self.handle_reversal(price, position_units, false);
                    return;
                }

                if self.position_size == 0.0 {
                    // Open short position
                    self.position_size = -position_units;
                    self.position_entry_price = price;
                    self.cash += position_units * price;
                } else if self.position_size > 0.0 {
                    // Close long position
                    self.cash += self.position_size * price;
                    self.position_size = 0.0;
                    self.position_entry_price = 0.0;
                } else {
                    // Already short, do nothing
                }
            }
            2 => {
                // HOLD - No action
            }
            _ => {
                // Invalid action - treat as HOLD
                warn!("Invalid PPO action index: {}. Treating as HOLD.", action_index);
            }
        }
    }

    /// Handle position reversal (2-phase process)
    ///
    /// This handles the case where an action would flip the position sign
    /// (e.g., long -> short or short -> long). The reversal is executed in two phases:
    /// 1. Close current position (realize P&L)
    /// 2. Open new position in opposite direction (may be partial if insufficient cash)
    ///
    /// # Arguments
    ///
    /// * `price` - Current market price
    /// * `position_units` - Total units requested (includes both close and open)
    /// * `is_buy` - True if BUY action (short->long), False if SELL action (long->short)
    #[allow(clippy::cognitive_complexity)]
    fn handle_reversal(&mut self, price: f32, position_units: f32, is_buy: bool) {
        // Phase 1: Close current position
        let phase1_units = self.position_size.abs();

        if self.position_size > 0.0 {
            // Closing long position: Sell generates cash
            self.cash += self.position_size * price;
        } else {
            // Closing short position: Buy to cover costs cash
            self.cash -= self.position_size.abs() * price;
        }

        let old_position = self.position_size;
        self.position_size = 0.0;
        self.position_entry_price = 0.0;

        warn!(
            "P2-R: Phase 1 complete - Closed position {:.2} \u{2192} 0.0, Cash: ${:.2}",
            old_position, self.cash
        );

        // Phase 2: Open new position in opposite direction
        let phase2_units = position_units - phase1_units;

        if phase2_units <= 0.0 {
            // No units remaining after closing - stay flat
            return;
        }

        // Calculate affordable position respecting cash reserve
        let portfolio_value = self.get_portfolio_value(price);
        let reserve_required = if self.cash_reserve_percent > 0.0 {
            portfolio_value * (self.cash_reserve_percent / 100.0)
        } else {
            0.0
        };

        if is_buy {
            // Opening long position after closing short
            let affordable_cash = (self.cash - reserve_required).max(0.0);
            let affordable_contracts = (affordable_cash / price).floor();
            let actual_contracts = affordable_contracts.min(phase2_units);

            if actual_contracts <= 0.0 {
                warn!(
                    "P2-R: Phase 2 SKIPPED - insufficient cash. Cash: ${:.2}, Reserve: ${:.2}, Target: {:.2}",
                    self.cash, reserve_required, phase2_units
                );
                return;
            }

            self.position_size = actual_contracts;
            self.position_entry_price = price;
            self.cash -= actual_contracts * price;

            warn!(
                "P2-R: Phase 2 complete - Opened long position {:.2} (target: {:.2}), Cash: ${:.2}",
                actual_contracts, phase2_units, self.cash
            );
        } else {
            // Opening short position after closing long (no cash constraint for shorts)
            self.position_size = -phase2_units;
            self.position_entry_price = price;
            self.cash += phase2_units * price;

            warn!(
                "P2-R: Phase 2 complete - Opened short position {:.2}, Cash: ${:.2}",
                -phase2_units, self.cash
            );
        }
    }

    /// Calculate current portfolio value (cash + unrealized P&L)
    ///
    /// # Arguments
    ///
    /// * `current_price` - Current market price
    ///
    /// # Returns
    ///
    /// Total portfolio value including unrealized P&L from open positions
    fn get_portfolio_value(&self, current_price: f32) -> f32 {
        // Portfolio value = cash + position_value_at_current_price
        // This naturally handles both long and short positions:
        // - Long: cash decreases when buying, position value increases with price
        // - Short: cash increases when selling, position value is negative (liability)
        self.cash + (self.position_size * current_price)
    }

    /// Reset portfolio to initial state (for new episode/epoch)
    ///
    /// This resets:
    /// - Cash to initial capital
    /// - Position size to 0 (flat)
    /// - Entry price to 0
    ///
    /// # Example
    ///
    /// ```
    /// use ml::ppo::portfolio_tracker::PortfolioTracker;
    ///
    /// let mut tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);
    /// tracker.execute_ppo_action(0, 100.0, 10.0);
    /// tracker.reset();
    /// assert_eq!(tracker.cash_balance(), 10_000.0);
    /// assert_eq!(tracker.current_position(), 0.0);
    /// ```
    pub const fn reset(&mut self) {
        self.cash = self.initial_capital;
        self.position_size = 0.0;
        self.position_entry_price = 0.0;
        self.last_price = 0.0;
        self.cumulative_transaction_costs = 0.0;
    }

    // ========== Public Accessors ==========

    /// Get current cash balance
    ///
    /// # Returns
    ///
    /// Current cash balance (may be negative if leveraged)
    pub const fn cash_balance(&self) -> f32 {
        self.cash
    }

    /// Get current position size
    ///
    /// # Returns
    ///
    /// Current position size (positive = long, negative = short, 0 = flat)
    pub const fn current_position(&self) -> f32 {
        self.position_size
    }

    /// Get total portfolio value (cash + unrealized P&L)
    ///
    /// This requires the current market price to calculate unrealized P&L.
    /// For flat positions, this is equivalent to cash balance.
    ///
    /// # Arguments
    ///
    /// * `current_price` - Current market price
    ///
    /// # Returns
    ///
    /// Total portfolio value including unrealized P&L
    pub fn total_value(&self, current_price: f32) -> f32 {
        self.get_portfolio_value(current_price)
    }

    /// Get average entry price for current position
    ///
    /// # Returns
    ///
    /// Average entry price (0.0 if no position)
    pub const fn average_entry_price(&self) -> f32 {
        self.position_entry_price
    }

    /// Get realized P&L
    ///
    /// Realized P&L is the profit/loss from closed positions.
    /// This is calculated as: (current cash - initial capital)
    ///
    /// # Returns
    ///
    /// Realized profit/loss
    pub fn realized_pnl(&self) -> f32 {
        self.cash - self.initial_capital
    }

    /// Get unrealized P&L
    ///
    /// Unrealized P&L is the profit/loss from open positions at current market price.
    ///
    /// # Arguments
    ///
    /// * `current_price` - Current market price
    ///
    /// # Returns
    ///
    /// Unrealized profit/loss (0.0 if no position)
    pub fn unrealized_pnl(&self, current_price: f32) -> f32 {
        // Unrealized P&L = current_portfolio_value - initial_capital
        // This is the profit/loss from open positions
        self.get_portfolio_value(current_price) - self.initial_capital
    }

    /// Get cumulative transaction costs
    ///
    /// # Returns
    ///
    /// Total transaction costs incurred across all trades
    pub const fn transaction_costs(&self) -> f32 {
        self.cumulative_transaction_costs
    }
}

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

    #[test]
    fn test_portfolio_tracker_initial_state() {
        let tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);

        assert_eq!(tracker.cash_balance(), 10_000.0);
        assert_eq!(tracker.current_position(), 0.0);
        assert_eq!(tracker.total_value(100.0), 10_000.0);
    }

    #[test]
    fn test_portfolio_tracker_buy_action() {
        let mut tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);
        tracker.execute_ppo_action(0, 100.0, 10.0);

        assert_eq!(tracker.current_position(), 10.0);
        assert_eq!(tracker.average_entry_price(), 100.0);
        assert_eq!(tracker.cash_balance(), 9_000.0);
    }

    #[test]
    fn test_portfolio_tracker_sell_action() {
        let mut tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);
        tracker.execute_ppo_action(1, 100.0, 10.0);

        assert_eq!(tracker.current_position(), -10.0);
        assert_eq!(tracker.average_entry_price(), 100.0);
        assert_eq!(tracker.cash_balance(), 11_000.0);
    }

    #[test]
    fn test_portfolio_tracker_hold_action() {
        let mut tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);
        let initial_cash = tracker.cash_balance();
        let initial_position = tracker.current_position();

        tracker.execute_ppo_action(2, 100.0, 10.0);

        assert_eq!(tracker.cash_balance(), initial_cash);
        assert_eq!(tracker.current_position(), initial_position);
    }

    #[test]
    fn test_portfolio_tracker_pnl_calculation_long() {
        let mut tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);
        tracker.execute_ppo_action(0, 100.0, 10.0);

        // Price rises to 110 -> +$100 unrealized profit
        let unrealized = tracker.unrealized_pnl(110.0);
        assert_eq!(unrealized, 100.0);

        let total_value = tracker.total_value(110.0);
        assert_eq!(total_value, 10_100.0);
    }

    #[test]
    fn test_portfolio_tracker_pnl_calculation_short() {
        let mut tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);
        tracker.execute_ppo_action(1, 100.0, 10.0);

        // Price falls to 90 -> +$100 unrealized profit (profitable for short)
        let unrealized = tracker.unrealized_pnl(90.0);
        assert_eq!(unrealized, 100.0);

        let total_value = tracker.total_value(90.0);
        assert_eq!(total_value, 10_100.0);
    }

    #[test]
    fn test_portfolio_tracker_close_long_position() {
        let mut tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);
        tracker.execute_ppo_action(0, 100.0, 10.0);
        tracker.execute_ppo_action(1, 110.0, 10.0);

        assert_eq!(tracker.current_position(), 0.0);
        assert_eq!(tracker.cash_balance(), 10_100.0);
    }

    #[test]
    fn test_portfolio_tracker_close_short_position() {
        let mut tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);
        tracker.execute_ppo_action(1, 100.0, 10.0);
        tracker.execute_ppo_action(0, 90.0, 10.0);

        assert_eq!(tracker.current_position(), 0.0);
        assert_eq!(tracker.cash_balance(), 10_100.0);
    }

    #[test]
    fn test_portfolio_tracker_reset() {
        let mut tracker = PortfolioTracker::new(10_000.0, 0.0001, 0.0);
        tracker.execute_ppo_action(0, 100.0, 10.0);
        tracker.reset();

        assert_eq!(tracker.cash_balance(), 10_000.0);
        assert_eq!(tracker.current_position(), 0.0);
        assert_eq!(tracker.average_entry_price(), 0.0);
    }
}
