foxhunt/crates/ml/src/ppo/ppo.rs

//! ACTUAL Working Proximal Policy Optimization (PPO) Implementation
//!
//! This module provides a complete, working PPO implementation with:
//! - Actor-Critic architecture with separate policy and value networks
//! - Real mathematical operations using candle-core v0.9.1
//! - Clipped surrogate objective function
//! - Generalized Advantage Estimation (GAE)
//! - Mini-batch SGD training with multiple epochs
//! - NO productions, todo!(), or unimplemented!() macros

#![allow(unsafe_code)] // Required for memory-mapped checkpoint loading

use candle_core::{DType, Device, Tensor};
use candle_nn::Module;
use candle_nn::{linear, Linear, Optimizer, VarBuilder, VarMap};
use candle_optimisers::adam::Adam;
use candle_optimisers::adam::ParamsAdam;
use rand::{thread_rng, Rng};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use tracing::{debug, info, warn};

use crate::gradient_accumulation::{accumulate_grads, check_gradients_finite, scale_grads};
use crate::tensor_ops::TensorOps;

use super::gae::GAEConfig;
use super::hidden_state_manager::HiddenStateManager;
use super::lstm_networks::{LSTMPolicyNetwork, LSTMValueNetwork};
use super::trajectories::{TrajectoryBatch, TrajectoryTensors};
use crate::dqn::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig};
use crate::dqn::portfolio_tracker::PortfolioTracker;
use crate::dqn::xavier_init::linear_xavier;
use crate::dqn::reward::RewardNormalizer;
use crate::dqn::TradingAction;
use crate::MLError;

/// Actor network variants supporting both MLP and LSTM architectures
///
/// This enum enables zero-cost runtime polymorphism for conditional LSTM usage.
/// Pattern matching on enum variants compiles to efficient jump tables with no
/// dynamic dispatch overhead (critical for HFT low-latency requirements).
#[allow(missing_debug_implementations)]
pub enum ActorNetwork {
    /// Standard feedforward Multi-Layer Perceptron (MLP) policy network
    ///
    /// Used when `PPOConfig::use_lstm = false` (default, backward compatible).
    /// Forward pass: state → logits (no hidden state propagation).
    MLP(PolicyNetwork),

    /// LSTM-augmented policy network for temporal dependencies
    ///
    /// Used when `PPOConfig::use_lstm = true`.
    /// Forward pass: (state, h_t, c_t) → (logits, new_h, new_c).
    /// Requires `HiddenStateManager` for state persistence across timesteps.
    LSTM(LSTMPolicyNetwork),
}

impl ActorNetwork {
    /// Get device for tensor operations
    pub fn device(&self) -> &Device {
        match self {
            ActorNetwork::MLP(network) => network.device(),
            ActorNetwork::LSTM(network) => network.device(),
        }
    }

    /// Get network variables for optimizer
    pub fn vars(&self) -> &VarMap {
        match self {
            ActorNetwork::MLP(network) => network.vars(),
            ActorNetwork::LSTM(network) => network.vars(),
        }
    }

    /// Forward pass (MLP-only method, requires explicit hidden state handling for LSTM)
    ///
    /// **WARNING**: This method only works for MLP networks. LSTM networks require
    /// explicit hidden state propagation via the training loop. Use `match` statements
    /// to handle both variants correctly.
    pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
        match self {
            ActorNetwork::MLP(network) => network.forward(input),
            ActorNetwork::LSTM(_) => Err(MLError::ModelError(
                "LSTM forward pass requires hidden states (h_t, c_t). Use match statements in training loop.".to_owned()
            )),
        }
    }

    /// Get action probabilities (MLP-only, softmax of logits)
    pub fn action_probabilities(&self, input: &Tensor) -> Result<Tensor, MLError> {
        match self {
            ActorNetwork::MLP(network) => network.action_probabilities(input),
            ActorNetwork::LSTM(_) => Err(MLError::ModelError(
                "LSTM action probabilities require hidden states. Use match statements in training loop.".to_owned()
            )),
        }
    }

    /// Sample action from policy (MLP-only)
    pub fn sample_action(&self, input: &Tensor) -> Result<(TradingAction, f32), MLError> {
        match self {
            ActorNetwork::MLP(network) => network.sample_action(input),
            ActorNetwork::LSTM(_) => Err(MLError::ModelError(
                "LSTM sample_action requires hidden states. Use match statements in training loop.".to_owned()
            )),
        }
    }

    /// Compute log probabilities for given actions (MLP-only)
    pub fn log_probs(&self, states: &Tensor, actions: &Tensor) -> Result<Tensor, MLError> {
        match self {
            ActorNetwork::MLP(network) => network.log_probs(states, actions),
            ActorNetwork::LSTM(_) => Err(MLError::ModelError(
                "LSTM log_probs requires hidden states. Use match statements in training loop.".to_owned()
            )),
        }
    }

    /// Compute entropy of action distribution (MLP-only)
    pub fn entropy(&self, states: &Tensor) -> Result<Tensor, MLError> {
        match self {
            ActorNetwork::MLP(network) => network.entropy(states),
            ActorNetwork::LSTM(_) => Err(MLError::ModelError(
                "LSTM entropy requires hidden states. Use match statements in training loop.".to_owned()
            )),
        }
    }
}

/// Critic network variants supporting both MLP and LSTM architectures
///
/// This enum enables zero-cost runtime polymorphism for conditional LSTM usage.
/// Pattern matching on enum variants compiles to efficient jump tables with no
/// dynamic dispatch overhead (critical for HFT low-latency requirements).
#[allow(missing_debug_implementations)]
pub enum CriticNetwork {
    /// Standard feedforward Multi-Layer Perceptron (MLP) value network
    ///
    /// Used when `PPOConfig::use_lstm = false` (default, backward compatible).
    /// Forward pass: state → value (no hidden state propagation).
    MLP(ValueNetwork),

    /// LSTM-augmented value network for temporal dependencies
    ///
    /// Used when `PPOConfig::use_lstm = true`.
    /// Forward pass: (state, h_t, c_t) → (value, new_h, new_c).
    /// Requires `HiddenStateManager` for state persistence across timesteps.
    LSTM(LSTMValueNetwork),
}

impl CriticNetwork {
    /// Get device for tensor operations
    pub fn device(&self) -> &Device {
        match self {
            CriticNetwork::MLP(network) => network.device(),
            CriticNetwork::LSTM(network) => network.device(),
        }
    }

    /// Get network variables for optimizer
    pub fn vars(&self) -> &VarMap {
        match self {
            CriticNetwork::MLP(network) => network.vars(),
            CriticNetwork::LSTM(network) => network.vars(),
        }
    }

    /// Forward pass (MLP-only method, requires explicit hidden state handling for LSTM)
    ///
    /// **WARNING**: This method only works for MLP networks. LSTM networks require
    /// explicit hidden state propagation via the training loop. Use `match` statements
    /// to handle both variants correctly.
    pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
        match self {
            CriticNetwork::MLP(network) => network.forward(input),
            CriticNetwork::LSTM(_) => Err(MLError::ModelError(
                "LSTM forward pass requires hidden states (h_t, c_t). Use match statements in training loop.".to_owned()
            )),
        }
    }
}

/// Configuration for `PPO` algorithm
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PPOConfig {
    /// State dimension
    pub state_dim: usize,
    /// Number of actions
    pub num_actions: usize,
    /// Policy network hidden dimensions
    pub policy_hidden_dims: Vec<usize>,
    /// Value network hidden dimensions
    pub value_hidden_dims: Vec<usize>,
    /// Learning rates
    pub policy_learning_rate: f64,
    pub value_learning_rate: f64,
    /// `PPO` clip parameter (epsilon)
    pub clip_epsilon: f32,
    /// Value function loss coefficient
    pub value_loss_coeff: f32,
    /// Entropy coefficient for exploration
    pub entropy_coeff: f32,
    /// GAE configuration
    pub gae_config: GAEConfig,
    /// Training parameters
    pub batch_size: usize,
    pub mini_batch_size: usize,
    pub num_epochs: usize,
    /// Maximum gradient norm for clipping
    pub max_grad_norm: f32,
    /// Early stopping enabled flag
    pub early_stopping_enabled: bool,
    /// Early stopping patience (epochs without improvement)
    pub early_stopping_patience: usize,
    /// Early stopping threshold (minimum improvement)
    pub early_stopping_min_delta: f64,
    /// Minimum epochs before early stopping can trigger
    pub early_stopping_min_epochs: usize,
    /// Maximum absolute position size (default: 2.0)
    pub max_position_absolute: f64,
    /// Transaction cost in basis points (default: 0.10%)
    pub transaction_cost_bps: f64,
    /// Cash reserve requirement as percentage (default: 20%)
    pub cash_reserve_pct: f64,
    /// Circuit breaker failure threshold (default: 5)
    pub circuit_breaker_threshold: usize,
    /// Use LSTM layers for temporal modeling (default: false for backward compatibility)
    pub use_lstm: bool,
    /// LSTM hidden dimension (default: 128)
    pub lstm_hidden_dim: usize,
    /// LSTM number of layers (default: 1)
    pub lstm_num_layers: usize,
    /// LSTM sequence length for training (default: 32)
    pub lstm_sequence_length: usize,
    /// Number of gradient accumulation steps (default: 1 = no accumulation)
    /// effective_batch = mini_batch_size * accumulation_steps
    pub accumulation_steps: usize,
    /// Optional higher clip bound for asymmetric clipping (clip-higher).
    /// When Some, uses clip(ratio, 1-clip_epsilon, 1+clip_epsilon_high).
    /// Prevents entropy collapse during long training. None = symmetric (default).
    pub clip_epsilon_high: Option<f32>,
}

impl Default for PPOConfig {
    fn default() -> Self {
        Self {
            state_dim: 64,
            num_actions: 45, // 5×3×3 factored action space (size × order type × duration)
            policy_hidden_dims: vec![128, 64],
            value_hidden_dims: vec![256, 128, 64], // Deeper network for better value approximation
            policy_learning_rate: 3e-5, // Reduced from 3e-4 to prevent gradient explosion
            value_learning_rate: 1e-4,  // Increased from 3e-5 to allow faster critic convergence
            clip_epsilon: 0.2,
            value_loss_coeff: 1.0, // Increased from 0.5 to prioritize value learning
            entropy_coeff: 0.05,   // Increased from 0.01 to encourage exploration
            gae_config: GAEConfig::default(),
            batch_size: 2048,
            mini_batch_size: 512, // Increased from 64 to prevent value network failure (88% gradient variance reduction)
            num_epochs: 20,       // Increased from 10 to allow critic to better fit value targets
            max_grad_norm: 0.5,
            early_stopping_enabled: true,
            early_stopping_patience: 10,
            early_stopping_min_delta: 1e-4,
            early_stopping_min_epochs: 20,
            max_position_absolute: 2.0,
            transaction_cost_bps: 0.10,
            cash_reserve_pct: 20.0,
            circuit_breaker_threshold: 5,
            use_lstm: false, // Backward compatible: standard MLP networks by default
            lstm_hidden_dim: 128,
            lstm_num_layers: 1,
            lstm_sequence_length: 32,
            accumulation_steps: 1,
            clip_epsilon_high: None,
        }
    }
}

/// Policy network for action probability distribution
#[allow(missing_debug_implementations)]
pub struct PolicyNetwork {
    layers: Vec<Linear>,
    device: Device,
    vars: VarMap,
}

impl PolicyNetwork {
    /// Create new policy network
    pub fn new(
        input_dim: usize,
        hidden_dims: &[usize],
        output_dim: usize,
        device: Device,
    ) -> Result<Self, MLError> {
        let vars = VarMap::new();
        let var_builder = VarBuilder::from_varmap(&vars, DType::F32, &device);

        let mut layers = Vec::new();
        let mut current_dim = input_dim;

        // Hidden layers
        for (i, &hidden_dim) in hidden_dims.iter().enumerate() {
            let layer = linear(
                current_dim,
                hidden_dim,
                var_builder.pp(format!("policy_layer_{}", i)),
            )
            .map_err(|e| {
                MLError::ModelError(format!("Failed to create policy layer {}: {}", i, e))
            })?;

            layers.push(layer);
            current_dim = hidden_dim;
        }

        // Output layer (logits for softmax)
        let output_layer = linear(current_dim, output_dim, var_builder.pp("policy_output"))
            .map_err(|e| {
                MLError::ModelError(format!("Failed to create policy output layer: {}", e))
            })?;

        layers.push(output_layer);

        Ok(Self {
            layers,
            device,
            vars,
        })
    }

    /// Load actor network from safetensors checkpoint via VarBuilder
    ///
    /// # Arguments
    /// * `vb` - VarBuilder loaded from safetensors file
    /// * `input_dim` - State dimension (must match training config)
    /// * `hidden_dims` - Hidden layer dimensions (must match training config)
    /// * `output_dim` - Number of actions (must match training config)
    /// * `device` - Device to load model on (CPU or CUDA)
    ///
    /// # Expected Checkpoint Structure
    /// ```text
    /// policy_layer_0.weight: [hidden_dims[0], input_dim]
    /// policy_layer_0.bias: [hidden_dims[0]]
    /// policy_layer_1.weight: [hidden_dims[1], hidden_dims[0]]
    /// policy_layer_1.bias: [hidden_dims[1]]
    /// ...
    /// policy_output.weight: [num_actions, hidden_dims`[last]`]
    /// policy_output.bias: [num_actions]
    /// ```
    ///
    /// # Returns
    /// Loaded PolicyNetwork with restored weights, or error if:
    /// - Checkpoint structure doesn't match config (wrong layer names/dimensions)
    /// - Tensor shapes are incompatible
    /// - Safetensors file is corrupted
    pub fn from_varbuilder(
        vb: VarBuilder<'_>,
        input_dim: usize,
        hidden_dims: &[usize],
        output_dim: usize,
        device: Device,
    ) -> Result<Self, MLError> {
        let mut layers = Vec::new();
        let mut current_dim = input_dim;

        // Load hidden layers from checkpoint
        for (i, &hidden_dim) in hidden_dims.iter().enumerate() {
            let layer_name = format!("policy_layer_{}", i);

            // Load weights and bias from checkpoint
            let layer = linear(current_dim, hidden_dim, vb.pp(&layer_name)).map_err(|e| {
                MLError::ModelError(format!(
                    "Failed to load actor layer {} from checkpoint: {}. \
                    Expected shape [{}, {}] for weights, got error: {}",
                    i, layer_name, hidden_dim, current_dim, e
                ))
            })?;

            layers.push(layer);
            current_dim = hidden_dim;
        }

        // Load output layer from checkpoint (action logits)
        let output_layer =
            linear(current_dim, output_dim, vb.pp("policy_output")).map_err(|e| {
                MLError::ModelError(format!(
                    "Failed to load actor output layer from checkpoint: {}. \
                    Expected shape [{}, {}] for weights",
                    e, output_dim, current_dim
                ))
            })?;

        layers.push(output_layer);

        // Create empty VarMap (weights are in VarBuilder, not VarMap for loaded models)
        let vars = VarMap::new();

        Ok(Self {
            layers,
            device,
            vars,
        })
    }

    /// Forward pass returning action logits
    pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
        let mut x = input.clone();

        // Pass through hidden layers with ReLU activation
        for (i, layer) in self.layers.iter().enumerate() {
            x = layer.forward(&x).map_err(|e| {
                MLError::ModelError(format!("Policy forward pass failed at layer {}: {}", i, e))
            })?;

            // Apply ReLU to all layers except the last
            if i < self.layers.len() - 1 {
                x = x
                    .relu()
                    .map_err(|e| MLError::ModelError(format!("ReLU activation failed: {}", e)))?;
            }
        }

        Ok(x)
    }

    /// Get action probabilities (softmax of logits)
    pub fn action_probabilities(&self, input: &Tensor) -> Result<Tensor, MLError> {
        let logits = self.forward(input)?;
        let probs = candle_nn::ops::softmax(&logits, candle_core::D::Minus1)
            .map_err(|e| MLError::ModelError(format!("Softmax failed: {}", e)))?;
        Ok(probs)
    }

    /// Sample action from policy
    pub fn sample_action(&self, input: &Tensor) -> Result<(TradingAction, f32), MLError> {
        let probs = self.action_probabilities(input)?;
        let probs_vec = probs
            .flatten_all()?
            .to_vec1::<f32>()
            .map_err(|e| MLError::ModelError(format!("Failed to extract probabilities: {}", e)))?;

        // Sample from categorical distribution
        let mut rng = thread_rng();
        let sample: f32 = rng.gen();
        let mut cumulative = 0.0;

        for (i, &prob) in probs_vec.iter().enumerate() {
            cumulative += prob;
            if sample <= cumulative {
                let action = TradingAction::from_int(i as u8)
                    .ok_or_else(|| MLError::InvalidInput(format!("Invalid action index: {}", i)))?;
                const EPSILON: f32 = 1e-8;
                let log_prob = (prob + EPSILON).ln();
                return Ok((action, log_prob));
            }
        }

        // Fallback to last action if rounding errors occur
        let last_idx = probs_vec.len() - 1;
        let action = TradingAction::from_int(last_idx as u8)
            .ok_or_else(|| MLError::InvalidInput(format!("Invalid action index: {}", last_idx)))?;
        const EPSILON: f32 = 1e-8;
        let log_prob = (probs_vec[last_idx] + EPSILON).ln();
        Ok((action, log_prob))
    }

    /// Compute log probabilities for given actions
    pub fn log_probs(&self, states: &Tensor, actions: &Tensor) -> Result<Tensor, MLError> {
        let logits = self.forward(states)?;
        let log_probs = candle_nn::ops::log_softmax(&logits, candle_core::D::Minus1)
            .map_err(|e| MLError::ModelError(format!("Log softmax failed: {}", e)))?;

        // Gather log probabilities for taken actions
        let actions_unsqueezed = actions.unsqueeze(1)?;
        let selected_log_probs = log_probs.gather(&actions_unsqueezed, 1)?.squeeze(1)?;

        Ok(selected_log_probs)
    }

    /// Compute entropy of action distribution
    pub fn entropy(&self, states: &Tensor) -> Result<Tensor, MLError> {
        let probs = self.action_probabilities(states)?;
        let log_probs =
            candle_nn::ops::log_softmax(&self.forward(states)?, candle_core::D::Minus1)?;

        // Entropy = -sum(p * log(p))
        let entropy_inner = (probs * log_probs)?.sum(candle_core::D::Minus1)?;
        let entropy = TensorOps::negate(&entropy_inner)?;
        Ok(entropy)
    }

    /// Get network variables
    pub fn vars(&self) -> &VarMap {
        &self.vars
    }

    /// Get device
    pub fn device(&self) -> &Device {
        &self.device
    }
}

/// Value network for state value estimation
#[allow(missing_debug_implementations)]
pub struct ValueNetwork {
    layers: Vec<Linear>,
    device: Device,
    vars: VarMap,
}

impl ValueNetwork {
    /// Create new value network
    pub fn new(input_dim: usize, hidden_dims: &[usize], device: Device) -> Result<Self, MLError> {
        let vars = VarMap::new();
        let var_builder = VarBuilder::from_varmap(&vars, DType::F32, &device);

        let mut layers = Vec::new();
        let mut current_dim = input_dim;

        // Hidden layers with Xavier initialization
        for (i, &hidden_dim) in hidden_dims.iter().enumerate() {
            let layer = linear_xavier(
                current_dim,
                hidden_dim,
                var_builder.pp(format!("value_layer_{}", i)),
            )
            .map_err(|e| {
                MLError::ModelError(format!("Failed to create value layer {}: {}", i, e))
            })?;

            layers.push(layer);
            current_dim = hidden_dim;
        }

        // Output layer (single value) with Xavier initialization
        let output_layer =
            linear_xavier(current_dim, 1, var_builder.pp("value_output")).map_err(|e| {
                MLError::ModelError(format!("Failed to create value output layer: {}", e))
            })?;

        layers.push(output_layer);

        Ok(Self {
            layers,
            device,
            vars,
        })
    }

    /// Load critic network from safetensors checkpoint via VarBuilder
    ///
    /// # Arguments
    /// * `vb` - VarBuilder loaded from safetensors file
    /// * `input_dim` - State dimension (must match training config)
    /// * `hidden_dims` - Hidden layer dimensions (must match training config)
    /// * `device` - Device to load model on (CPU or CUDA)
    ///
    /// # Expected Checkpoint Structure
    /// ```text
    /// value_layer_0.weight: [hidden_dims[0], input_dim]
    /// value_layer_0.bias: [hidden_dims[0]]
    /// value_layer_1.weight: [hidden_dims[1], hidden_dims[0]]
    /// value_layer_1.bias: [hidden_dims[1]]
    /// ...
    /// value_output.weight: [1, hidden_dims`[last]`]
    /// value_output.bias: `[1]`
    /// ```
    ///
    /// # Returns
    /// Loaded ValueNetwork with restored weights, or error if:
    /// - Checkpoint structure doesn't match config (wrong layer names/dimensions)
    /// - Tensor shapes are incompatible
    /// - Safetensors file is corrupted
    pub fn from_varbuilder(
        vb: VarBuilder<'_>,
        input_dim: usize,
        hidden_dims: &[usize],
        device: Device,
    ) -> Result<Self, MLError> {
        let mut layers = Vec::new();
        let mut current_dim = input_dim;

        // Load hidden layers from checkpoint
        for (i, &hidden_dim) in hidden_dims.iter().enumerate() {
            let layer_name = format!("value_layer_{}", i);

            // Load weights and bias from checkpoint
            let layer = linear(current_dim, hidden_dim, vb.pp(&layer_name)).map_err(|e| {
                MLError::ModelError(format!(
                    "Failed to load critic layer {} from checkpoint: {}. \
                    Expected shape [{}, {}] for weights, got error: {}",
                    i, layer_name, hidden_dim, current_dim, e
                ))
            })?;

            layers.push(layer);
            current_dim = hidden_dim;
        }

        // Load output layer from checkpoint (single value output)
        let output_layer = linear(current_dim, 1, vb.pp("value_output")).map_err(|e| {
            MLError::ModelError(format!(
                "Failed to load critic output layer from checkpoint: {}. \
                    Expected shape [1, {}] for weights",
                e, current_dim
            ))
        })?;

        layers.push(output_layer);

        // Create empty VarMap (weights are in VarBuilder, not VarMap for loaded models)
        let vars = VarMap::new();

        Ok(Self {
            layers,
            device,
            vars,
        })
    }

    /// Forward pass returning state values
    pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
        let mut x = input.clone();

        // Pass through hidden layers with ReLU activation
        for (i, layer) in self.layers.iter().enumerate() {
            x = layer.forward(&x).map_err(|e| {
                MLError::ModelError(format!("Value forward pass failed at layer {}: {}", i, e))
            })?;

            // Apply ReLU to all layers except the last
            if i < self.layers.len() - 1 {
                x = x
                    .relu()
                    .map_err(|e| MLError::ModelError(format!("ReLU activation failed: {}", e)))?;
            }
        }

        // Squeeze the last dimension (from [batch, 1] to [batch])
        x = x.squeeze(1)?;

        Ok(x)
    }

    /// Get network variables
    pub fn vars(&self) -> &VarMap {
        &self.vars
    }

    /// Get device
    pub fn device(&self) -> &Device {
        &self.device
    }
}

/// Working `PPO` implementation with support for both MLP and LSTM architectures
///
/// # Architecture Modes
///
/// ## MLP Mode (default, `use_lstm = false`)
/// - `actor`: ActorNetwork::MLP(PolicyNetwork)
/// - `critic`: CriticNetwork::MLP(ValueNetwork)
/// - `hidden_state_manager`: None
/// - Forward pass: state → action/value (no temporal dependencies)
///
/// ## LSTM Mode (`use_lstm = true`)
/// - `actor`: ActorNetwork::LSTM(LSTMPolicyNetwork)
/// - `critic`: CriticNetwork::LSTM(LSTMValueNetwork)
/// - `hidden_state_manager`: Some(HiddenStateManager)
/// - Forward pass: (state, h_t, c_t) → (action/value, new_h, new_c)
/// - Requires sequence batching via `TrajectoryBatch::to_sequences()`
#[allow(missing_debug_implementations)]
pub struct PPO {
    /// `PPO` configuration
    config: PPOConfig,
    /// Policy network (actor) - supports both MLP and LSTM variants
    pub actor: ActorNetwork,
    /// Value network (critic) - supports both MLP and LSTM variants
    pub critic: CriticNetwork,
    /// Policy optimizer
    policy_optimizer: Option<Adam>,
    /// Value optimizer
    value_optimizer: Option<Adam>,
    /// Training step counter
    pub training_steps: u64,
    /// Portfolio tracker for cash, position, and P&L management
    pub portfolio_tracker: PortfolioTracker,
    /// Reward normalizer for ~N(0,1) distribution
    pub reward_normalizer: Option<RewardNormalizer>,
    /// Circuit breaker for failure detection and cooldown
    pub circuit_breaker: Option<CircuitBreaker>,
    /// Transaction cost in basis points
    pub transaction_cost_bps: Option<f64>,
    /// Maximum absolute position size
    pub max_position_absolute: Option<f64>,
    /// Hidden state manager for LSTM (None if use_lstm = false)
    pub hidden_state_manager: Option<HiddenStateManager>,
}

/// Backward-compatibility alias: `WorkingPPO` is now [`PPO`].
pub type WorkingPPO = PPO;

impl PPO {
    /// Create new `PPO` with GPU by default (falls back to CPU if unavailable)
    pub fn new(config: PPOConfig) -> Result<Self, MLError> {
        let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
        Self::with_device(config, device)
    }

    /// Create new working `PPO` with specified device (GPU or CPU)
    ///
    /// # Network Selection
    ///
    /// This constructor dynamically selects MLP or LSTM networks based on `config.use_lstm`:
    /// - `use_lstm = false` (default): Creates ActorNetwork::MLP + CriticNetwork::MLP
    /// - `use_lstm = true`: Creates ActorNetwork::LSTM + CriticNetwork::LSTM
    ///
    /// # LSTM Requirements
    ///
    /// When `use_lstm = true`, ensure:
    /// 1. `config.lstm_hidden_dim` is set (e.g., 128)
    /// 2. `config.lstm_num_layers` is set (e.g., 1-2)
    /// 3. Training uses `TrajectoryBatch::to_sequences()` for BPTT
    /// 4. `HiddenStateManager` is initialized automatically
    pub fn with_device(config: PPOConfig, device: Device) -> Result<Self, MLError> {
        if config.state_dim == 0 {
            return Err(MLError::ConfigError {
                reason: "PPO requires state_dim > 0".to_owned(),
            });
        }
        if config.num_actions == 0 {
            return Err(MLError::ConfigError {
                reason: "PPO requires num_actions > 0".to_owned(),
            });
        }

        // Extract values before moving config
        let use_lstm = config.use_lstm;
        let lstm_hidden_dim = config.lstm_hidden_dim;
        let lstm_num_layers = config.lstm_num_layers;
        let transaction_cost_bps = config.transaction_cost_bps;
        let max_position_absolute = config.max_position_absolute;

        // Create networks based on use_lstm flag
        let (actor, critic) = if use_lstm {
            // LSTM mode: Create LSTM-augmented networks
            let lstm_actor = LSTMPolicyNetwork::new(
                config.state_dim,
                lstm_hidden_dim,
                lstm_num_layers,
                config.num_actions,
                device.clone(),
            )?;

            let lstm_critic = LSTMValueNetwork::new(
                config.state_dim,
                lstm_hidden_dim,
                lstm_num_layers,
                device.clone(),
            )?;

            (
                ActorNetwork::LSTM(lstm_actor),
                CriticNetwork::LSTM(lstm_critic),
            )
        } else {
            // MLP mode (default): Create standard feedforward networks
            let mlp_actor = PolicyNetwork::new(
                config.state_dim,
                &config.policy_hidden_dims,
                config.num_actions,
                device.clone(),
            )?;

            let mlp_critic = ValueNetwork::new(
                config.state_dim,
                &config.value_hidden_dims,
                device.clone(),
            )?;

            (
                ActorNetwork::MLP(mlp_actor),
                CriticNetwork::MLP(mlp_critic),
            )
        };

        // Create portfolio tracker (initial capital: 10,000, spread: 0.0001, cash reserve: config%)
        let portfolio_tracker = PortfolioTracker::new(
            10_000.0,
            0.0001,
            config.cash_reserve_pct,
        );

        // Create reward normalizer (enabled by default)
        let reward_normalizer = Some(RewardNormalizer::new());

        // Create circuit breaker with configured threshold
        let circuit_breaker_config = CircuitBreakerConfig {
            failure_threshold: config.circuit_breaker_threshold,
            success_threshold: 3,
            timeout_duration: std::time::Duration::from_secs(60),
            half_open_max_calls: 2,
        };
        let circuit_breaker = Some(CircuitBreaker::new(circuit_breaker_config));

        // Initialize hidden state manager if LSTM is enabled
        let hidden_state_manager = use_lstm
            .then(|| {
                // Batch size of 1 for single environment (will be updated during training)
                HiddenStateManager::new(
                    lstm_num_layers,
                    1, // Default batch size, will be resized when needed
                    lstm_hidden_dim,
                    &critic.device(),
                )
            })
            .transpose()?;

        Ok(Self {
            config,
            actor,
            critic,
            policy_optimizer: None,
            value_optimizer: None,
            training_steps: 0,
            portfolio_tracker,
            reward_normalizer,
            circuit_breaker,
            transaction_cost_bps: Some(transaction_cost_bps),
            max_position_absolute: Some(max_position_absolute),
            hidden_state_manager,
        })
    }

    /// Select action and get value estimate
    pub fn act(&self, state: &[f32]) -> Result<(TradingAction, f32), MLError> {
        let (action, _log_prob, value) = self.act_with_log_prob(state)?;
        Ok((action, value))
    }

    /// Select action and return (action, log_prob, value).
    ///
    /// Returns the **real** policy log-probability needed for PPO importance
    /// sampling. Use this instead of [`act`] when collecting trajectories
    /// for training.
    pub fn act_with_log_prob(&self, state: &[f32]) -> Result<(TradingAction, f32, f32), MLError> {
        let state_tensor = Tensor::from_vec(
            state.to_vec(),
            (1, self.config.state_dim),
            self.actor.device(),
        )
        .map_err(|e| MLError::ModelError(format!("Failed to create state tensor: {}", e)))?;

        // Get action and log-probability from policy
        let (action, log_prob) = self.actor.sample_action(&state_tensor)?;

        // Get value estimate
        let value_tensor = self.critic.forward(&state_tensor)?;
        let value = value_tensor
            .get(0)
            .map_err(|e| MLError::ModelError(format!("Failed to get value element: {}", e)))?
            .to_scalar::<f32>()
            .map_err(|e| MLError::ModelError(format!("Failed to extract value: {}", e)))?;

        Ok((action, log_prob, value))
    }

    /// Update `PPO` networks with trajectory batch
    pub fn update(&mut self, batch: &mut TrajectoryBatch) -> Result<(f32, f32), MLError> {
        // Check circuit breaker before training
        if let Some(ref circuit_breaker) = self.circuit_breaker {
            if !circuit_breaker.allow_request() {
                warn!("Circuit breaker is open - skipping training update");
                return Ok((0.0, 0.0));
            }
        }

        // Initialize optimizers if not done
        self.init_optimizers()?;

        // Apply reward normalization if enabled
        if let Some(ref mut normalizer) = self.reward_normalizer {
            for reward in &batch.rewards {
                normalizer.update(*reward as f64);
            }
            // Normalize all rewards in batch
            let normalized_rewards: Vec<f32> = batch.rewards.iter()
                .map(|&r| {
                    let norm = normalizer.normalize(r as f64);
                    norm.clamp(-1.0, 1.0) as f32
                })
                .collect();
            batch.rewards = normalized_rewards;
        }

        // Normalize advantages
        batch.normalize_advantages()?;

        // Branch on network type for training
        match (&self.actor, &self.critic) {
            (ActorNetwork::MLP(_), CriticNetwork::MLP(_)) => {
                // Existing MLP training loop (keep as-is)
                self.update_mlp(batch)
            },
            (ActorNetwork::LSTM(_), CriticNetwork::LSTM(_)) => {
                // New LSTM training loop
                self.update_lstm(batch)
            },
            _ => {
                Err(MLError::ModelError(
                    "Actor and Critic must both be MLP or both be LSTM".to_owned()
                ))
            }
        }
    }

    /// MLP-specific update logic (original implementation)
    fn update_mlp(&mut self, batch: &mut TrajectoryBatch) -> Result<(f32, f32), MLError> {
        // Convert batch to tensors
        let device = self.actor.device();
        let _batch_tensors = batch.to_tensors(device, self.config.state_dim)?;

        let accumulation_steps = self.config.accumulation_steps.max(1);
        if accumulation_steps > 1 {
            info!(
                "Using gradient accumulation: {} steps (effective batch = {})",
                accumulation_steps,
                self.config.mini_batch_size * accumulation_steps
            );
        }

        let mut total_policy_loss = 0.0;
        let mut total_value_loss = 0.0;
        let mut num_updates = 0;

        // Train for multiple epochs
        for epoch in 0..self.config.num_epochs {
            // Create mini-batches
            let mini_batches = batch.create_mini_batches(self.config.mini_batch_size);

            let mut policy_grad_accumulator: Option<candle_core::backprop::GradStore> = None;
            let mut value_grad_accumulator: Option<candle_core::backprop::GradStore> = None;
            let mut accum_step: usize = 0;

            for mini_batch in mini_batches {
                let mini_tensors = mini_batch.to_tensors(device, self.config.state_dim)?;

                // Compute losses
                let policy_loss = self.compute_policy_loss(&mini_tensors)?;
                let value_loss = self.compute_value_loss(&mini_tensors)?;

                // Extract scalar values for NaN check
                let policy_loss_scalar = policy_loss.to_scalar::<f32>().map_err(|e| {
                    MLError::TrainingError(format!("Failed to extract policy loss: {}", e))
                })?;
                let value_loss_scalar = value_loss.to_scalar::<f32>().map_err(|e| {
                    MLError::TrainingError(format!("Failed to extract value loss: {}", e))
                })?;

                // NaN detection every 10 epochs
                if epoch % 10 == 0 {
                    if policy_loss_scalar.is_nan() {
                        return Err(MLError::TrainingError(
                            format!("NaN detected in policy loss at epoch {} - training unstable. \
                                    Consider reducing learning rate or increasing entropy coefficient.", epoch)
                        ));
                    }
                    if value_loss_scalar.is_nan() {
                        return Err(MLError::TrainingError(format!(
                            "NaN detected in value loss at epoch {} - training unstable. \
                                    Consider reducing learning rate.",
                            epoch
                        )));
                    }
                }

                // Accumulate policy gradients
                let actor_vars = self.actor.vars().all_vars();
                let policy_grads = policy_loss.backward().map_err(|e| {
                    MLError::TrainingError(format!("Policy backward failed: {}", e))
                })?;
                accumulate_grads(&mut policy_grad_accumulator, policy_grads, &actor_vars)?;

                // Accumulate value gradients
                let critic_vars = self.critic.vars().all_vars();
                let value_grads = value_loss.backward().map_err(|e| {
                    MLError::TrainingError(format!("Value backward failed: {}", e))
                })?;
                accumulate_grads(&mut value_grad_accumulator, value_grads, &critic_vars)?;

                accum_step += 1;

                // Step when we've accumulated enough
                if accum_step >= accumulation_steps {
                    // Scale and step policy optimizer
                    if let Some(ref mut policy_grads) = policy_grad_accumulator {
                        scale_grads(
                            policy_grads,
                            &actor_vars,
                            1.0 / accumulation_steps as f64,
                        )?;
                        check_gradients_finite(policy_grads, &actor_vars)
                            .map_err(|e| MLError::TrainingError(format!("Policy gradient NaN: {}", e)))?;

                        let policy_grad_norm =
                            Self::compute_gradient_norm(&actor_vars, policy_grads)?;

                        if policy_grad_norm > self.config.max_grad_norm as f64 {
                            warn!(
                                "Policy gradient norm {:.4} exceeds max {:.4}. Consider reducing learning rate.",
                                policy_grad_norm, self.config.max_grad_norm
                            );
                        }

                        debug!("Policy gradient norm: {:.4}", policy_grad_norm);

                        if let Some(ref mut optimizer) = self.policy_optimizer {
                            optimizer.step(policy_grads).map_err(|e| {
                                MLError::TrainingError(format!(
                                    "Policy optimizer step failed: {}",
                                    e
                                ))
                            })?;
                        }
                    }

                    // Scale and step value optimizer
                    if let Some(ref mut value_grads) = value_grad_accumulator {
                        scale_grads(
                            value_grads,
                            &critic_vars,
                            1.0 / accumulation_steps as f64,
                        )?;
                        check_gradients_finite(value_grads, &critic_vars)
                            .map_err(|e| MLError::TrainingError(format!("Value gradient NaN: {}", e)))?;

                        let value_grad_norm =
                            Self::compute_gradient_norm(&critic_vars, value_grads)?;

                        if value_grad_norm > self.config.max_grad_norm as f64 {
                            warn!(
                                "Value gradient norm {:.4} exceeds max {:.4}. Consider reducing learning rate.",
                                value_grad_norm, self.config.max_grad_norm
                            );
                        }

                        debug!("Value gradient norm: {:.4}", value_grad_norm);

                        if let Some(ref mut optimizer) = self.value_optimizer {
                            optimizer.step(value_grads).map_err(|e| {
                                MLError::TrainingError(format!(
                                    "Value optimizer step failed: {}",
                                    e
                                ))
                            })?;
                        }
                    }

                    // Reset accumulators
                    policy_grad_accumulator = None;
                    value_grad_accumulator = None;
                    accum_step = 0;
                }

                total_policy_loss += policy_loss_scalar;
                total_value_loss += value_loss_scalar;
                num_updates += 1;
            }

            // Handle remaining accumulated gradients (if mini_batches not evenly divisible)
            if accum_step > 0 {
                let actor_vars = self.actor.vars().all_vars();
                let critic_vars = self.critic.vars().all_vars();

                if let Some(ref mut policy_grads) = policy_grad_accumulator {
                    scale_grads(policy_grads, &actor_vars, 1.0 / accum_step as f64)?;
                    check_gradients_finite(policy_grads, &actor_vars)
                        .map_err(|e| MLError::TrainingError(format!("Policy gradient NaN: {}", e)))?;

                    let policy_grad_norm =
                        Self::compute_gradient_norm(&actor_vars, policy_grads)?;

                    if policy_grad_norm > self.config.max_grad_norm as f64 {
                        warn!(
                            "Policy gradient norm {:.4} exceeds max {:.4}. Consider reducing learning rate.",
                            policy_grad_norm, self.config.max_grad_norm
                        );
                    }

                    debug!(
                        "Policy gradient norm: {:.4} (remainder)",
                        policy_grad_norm
                    );

                    if let Some(ref mut optimizer) = self.policy_optimizer {
                        optimizer.step(policy_grads).map_err(|e| {
                            MLError::TrainingError(format!(
                                "Policy optimizer step failed: {}",
                                e
                            ))
                        })?;
                    }
                }

                if let Some(ref mut value_grads) = value_grad_accumulator {
                    scale_grads(value_grads, &critic_vars, 1.0 / accum_step as f64)?;
                    check_gradients_finite(value_grads, &critic_vars)
                        .map_err(|e| MLError::TrainingError(format!("Value gradient NaN: {}", e)))?;

                    let value_grad_norm =
                        Self::compute_gradient_norm(&critic_vars, value_grads)?;

                    if value_grad_norm > self.config.max_grad_norm as f64 {
                        warn!(
                            "Value gradient norm {:.4} exceeds max {:.4}. Consider reducing learning rate.",
                            value_grad_norm, self.config.max_grad_norm
                        );
                    }

                    debug!(
                        "Value gradient norm: {:.4} (remainder)",
                        value_grad_norm
                    );

                    if let Some(ref mut optimizer) = self.value_optimizer {
                        optimizer.step(value_grads).map_err(|e| {
                            MLError::TrainingError(format!(
                                "Value optimizer step failed: {}",
                                e
                            ))
                        })?;
                    }
                }
            }
        }

        self.training_steps += 1;

        let avg_policy_loss = total_policy_loss / num_updates as f32;
        let avg_value_loss = total_value_loss / num_updates as f32;

        // Update circuit breaker based on training success
        if let Some(ref circuit_breaker) = self.circuit_breaker {
            // Consider training successful if losses are reasonable (not NaN, not exploding)
            if avg_policy_loss.is_finite() && avg_value_loss.is_finite() {
                circuit_breaker.record_success();
            } else {
                circuit_breaker.record_failure();
                warn!("Training failure detected: policy_loss={}, value_loss={}", avg_policy_loss, avg_value_loss);
            }
        }

        Ok((avg_policy_loss, avg_value_loss))
    }

    /// LSTM-specific update logic with sequence processing and hidden state management
    fn update_lstm(&mut self, batch: &mut TrajectoryBatch) -> Result<(f32, f32), MLError> {
        // Verify hidden state manager exists
        let hidden_state_manager = self.hidden_state_manager.as_mut()
            .ok_or_else(|| MLError::ModelError(
                "LSTM mode requires HiddenStateManager but it's not initialized".to_owned()
            ))?;

        // Create sequences from batch
        let sequences = batch.to_sequences(self.config.lstm_sequence_length);

        let device = self.actor.device();
        let mut total_policy_loss = 0.0;
        let mut total_value_loss = 0.0;
        let mut num_updates = 0;

        // Extract LSTM networks (we already validated both are LSTM in match)
        let (actor_lstm, critic_lstm) = match (&self.actor, &self.critic) {
            (ActorNetwork::LSTM(actor), CriticNetwork::LSTM(critic)) => (actor, critic),
            _ => return Err(MLError::ConfigError { reason: "Expected both actor and critic to be LSTM".to_owned() }),
        };

        // Train for multiple epochs
        for epoch in 0..self.config.num_epochs {
            // Process each sequence
            for sequence in &sequences {
                let seq_len = sequence.length();

                // Reset hidden states at the start of each sequence
                hidden_state_manager.reset_all()?;

                // Collect outputs for the entire sequence
                let mut seq_log_probs = Vec::new();
                let mut seq_values = Vec::new();

                // Process sequence timestep-by-timestep
                for t in 0..seq_len {
                    // Convert single state to tensor [1, state_dim]
                    let state_tensor = Tensor::from_vec(
                        sequence.states[t].clone(),
                        (1, self.config.state_dim),
                        device,
                    ).map_err(|e| MLError::TensorOperationError(
                        format!("Failed to create state tensor: {}", e)
                    ))?;

                    // Get current hidden states
                    let (policy_h, policy_c) = hidden_state_manager.get_policy_state();
                    let (value_h, value_c) = hidden_state_manager.get_value_state();

                    // Actor forward pass with LSTM
                    let (logits, new_policy_h, new_policy_c) = actor_lstm.forward(
                        &state_tensor,
                        &policy_h,
                        &policy_c,
                    )?;

                    // Critic forward pass with LSTM
                    let (value, new_value_h, new_value_c) = critic_lstm.forward(
                        &state_tensor,
                        &value_h,
                        &value_c,
                    )?;

                    // Update hidden states
                    hidden_state_manager.update_policy_state(new_policy_h, new_policy_c)?;
                    hidden_state_manager.update_value_state(new_value_h, new_value_c)?;

                    // Compute log probability of taken action
                    let log_probs_dist = candle_nn::ops::log_softmax(&logits, candle_core::D::Minus1)
                        .map_err(|e| MLError::ModelError(format!("Log softmax failed: {}", e)))?;

                    let action_idx = sequence.actions[t].to_int() as i64;
                    let action_tensor = Tensor::from_vec(
                        vec![action_idx],
                        (1, 1),
                        device,
                    ).map_err(|e| MLError::TensorOperationError(
                        format!("Failed to create action tensor: {}", e)
                    ))?;

                    let log_prob = log_probs_dist.gather(&action_tensor, 1)?
                        .squeeze(1)?
                        .get(0)?
                        .to_scalar::<f32>()
                        .map_err(|e| MLError::TensorOperationError(
                            format!("Failed to extract log prob: {}", e)
                        ))?;

                    let value_scalar = value.get(0)?
                        .to_scalar::<f32>()
                        .map_err(|e| MLError::TensorOperationError(
                            format!("Failed to extract value: {}", e)
                        ))?;

                    seq_log_probs.push(log_prob);
                    seq_values.push(value_scalar);

                    // Reset hidden states on episode boundary
                    if sequence.dones[t] {
                        hidden_state_manager.reset_all()?;
                    }
                }

                // Convert sequence data to tensors for loss computation
                let seq_old_log_probs = Tensor::from_vec(
                    sequence.log_probs.clone(),
                    (seq_len,),
                    device,
                )?;

                let seq_new_log_probs = Tensor::from_vec(
                    seq_log_probs,
                    (seq_len,),
                    device,
                )?;

                let seq_advantages = Tensor::from_vec(
                    sequence.advantages.clone(),
                    (seq_len,),
                    device,
                )?;

                let seq_returns = Tensor::from_vec(
                    sequence.returns.clone(),
                    (seq_len,),
                    device,
                )?;

                let seq_values_tensor = Tensor::from_vec(
                    seq_values,
                    (seq_len,),
                    device,
                )?;

                // Compute policy loss (PPO clipped objective)
                let log_ratio = (&seq_new_log_probs - &seq_old_log_probs)?;
                let ratio = log_ratio.exp()?;

                let clip_epsilon_tensor = Tensor::from_vec(
                    vec![self.config.clip_epsilon; seq_len],
                    (seq_len,),
                    device,
                )?;

                let one_tensor = Tensor::ones((seq_len,), DType::F32, device)?;
                let clip_min = (&one_tensor - &clip_epsilon_tensor)?;
                let clip_max = if let Some(eps_high) = self.config.clip_epsilon_high {
                    let clip_high_tensor = Tensor::from_vec(
                        vec![eps_high; seq_len],
                        (seq_len,),
                        device,
                    ).map_err(|e| MLError::TrainingError(format!("Failed to create clip_high tensor: {}", e)))?;
                    (&one_tensor + &clip_high_tensor)?
                } else {
                    (&one_tensor + &clip_epsilon_tensor)?
                };

                let clipped_ratio = ratio.clamp(&clip_min, &clip_max)?;

                let surr1 = (&ratio * &seq_advantages)?;
                let surr2 = (&clipped_ratio * &seq_advantages)?;
                let policy_loss_raw = TensorOps::elementwise_min(&surr1, &surr2)?;

                // Compute entropy from log-probabilities: H = -mean(log_probs)
                // For a well-calibrated policy, entropy measures exploration breadth
                let entropy = seq_new_log_probs.neg()?.mean_all()?;
                let entropy_bonus = TensorOps::scalar_mul(&entropy, self.config.entropy_coeff as f64)?;

                let policy_loss_mean = policy_loss_raw.mean_all()?;
                let policy_loss_inner = (policy_loss_mean + entropy_bonus)?;
                let policy_loss = TensorOps::negate(&policy_loss_inner)?;

                // Compute value loss
                let value_loss = (&seq_values_tensor - &seq_returns)?
                    .powf(2.0)?
                    .mean_all()?;
                let scaled_value_loss = TensorOps::scalar_mul(&value_loss, self.config.value_loss_coeff as f64)?;

                // Extract scalar values for NaN check
                let policy_loss_scalar = policy_loss.to_scalar::<f32>().map_err(|e| {
                    MLError::TrainingError(format!("Failed to extract policy loss: {}", e))
                })?;
                let value_loss_scalar = scaled_value_loss.to_scalar::<f32>().map_err(|e| {
                    MLError::TrainingError(format!("Failed to extract value loss: {}", e))
                })?;

                // NaN detection every 10 epochs
                if epoch % 10 == 0 {
                    if policy_loss_scalar.is_nan() {
                        return Err(MLError::TrainingError(
                            format!("NaN detected in policy loss at epoch {} - LSTM training unstable. \
                                    Consider reducing learning rate or sequence length.", epoch)
                        ));
                    }
                    if value_loss_scalar.is_nan() {
                        return Err(MLError::TrainingError(format!(
                            "NaN detected in value loss at epoch {} - LSTM training unstable. \
                                    Consider reducing learning rate.",
                            epoch
                        )));
                    }
                }

                // Update policy network with gradient norm monitoring
                let actor_vars = self.actor.vars().all_vars();
                let policy_grads = policy_loss.backward().map_err(|e| {
                    MLError::TrainingError(format!("Policy backward failed: {}", e))
                })?;
                check_gradients_finite(&policy_grads, &actor_vars)
                    .map_err(|e| MLError::TrainingError(format!("Policy gradient NaN: {}", e)))?;

                let policy_grad_norm =
                    Self::compute_gradient_norm(&actor_vars, &policy_grads)?;

                if let Some(ref mut optimizer) = self.policy_optimizer {
                    optimizer.step(&policy_grads).map_err(|e| {
                        MLError::TrainingError(format!("Policy optimizer step failed: {}", e))
                    })?;
                }

                if policy_grad_norm > self.config.max_grad_norm as f64 {
                    warn!(
                        "LSTM policy gradient norm {:.4} exceeds max {:.4}. Consider reducing learning rate or sequence length.",
                        policy_grad_norm, self.config.max_grad_norm
                    );
                }

                debug!("LSTM policy gradient norm: {:.4}", policy_grad_norm);

                // Update value network with gradient norm monitoring
                let critic_vars = self.critic.vars().all_vars();
                let value_grads = scaled_value_loss.backward().map_err(|e| {
                    MLError::TrainingError(format!("Value backward failed: {}", e))
                })?;
                check_gradients_finite(&value_grads, &critic_vars)
                    .map_err(|e| MLError::TrainingError(format!("Value gradient NaN: {}", e)))?;

                let value_grad_norm =
                    Self::compute_gradient_norm(&critic_vars, &value_grads)?;

                if let Some(ref mut optimizer) = self.value_optimizer {
                    optimizer.step(&value_grads).map_err(|e| {
                        MLError::TrainingError(format!("Value optimizer step failed: {}", e))
                    })?;
                }

                if value_grad_norm > self.config.max_grad_norm as f64 {
                    warn!(
                        "LSTM value gradient norm {:.4} exceeds max {:.4}. Consider reducing learning rate.",
                        value_grad_norm, self.config.max_grad_norm
                    );
                }

                debug!("LSTM value gradient norm: {:.4}", value_grad_norm);

                total_policy_loss += policy_loss_scalar;
                total_value_loss += value_loss_scalar;
                num_updates += 1;
            }
        }

        self.training_steps += 1;

        let avg_policy_loss = total_policy_loss / num_updates as f32;
        let avg_value_loss = total_value_loss / num_updates as f32;

        // Update circuit breaker based on training success
        if let Some(ref circuit_breaker) = self.circuit_breaker {
            if avg_policy_loss.is_finite() && avg_value_loss.is_finite() {
                circuit_breaker.record_success();
            } else {
                circuit_breaker.record_failure();
                warn!("LSTM training failure detected: policy_loss={}, value_loss={}", avg_policy_loss, avg_value_loss);
            }
        }

        Ok((avg_policy_loss, avg_value_loss))
    }

    /// Compute the L2 norm of all gradients for monitoring
    ///
    /// This follows the same pattern as `ContinuousPPO::compute_gradient_norm`.
    /// The norm is used to detect gradient explosion; when it exceeds
    /// `max_grad_norm`, a warning is logged recommending learning rate reduction.
    fn compute_gradient_norm(
        vars: &[candle_core::Var],
        grads: &candle_core::backprop::GradStore,
    ) -> Result<f64, MLError> {
        let mut total_norm_sq = 0.0_f64;

        for var in vars {
            if let Some(grad) = grads.get(var) {
                let grad_norm_sq = grad
                    .sqr()
                    .map_err(|e| {
                        MLError::TrainingError(format!("Failed to square gradient: {}", e))
                    })?
                    .sum_all()
                    .map_err(|e| {
                        MLError::TrainingError(format!("Failed to sum gradient: {}", e))
                    })?
                    .to_vec0::<f32>()
                    .map_err(|e| {
                        MLError::TrainingError(format!(
                            "Failed to extract gradient norm: {}",
                            e
                        ))
                    })? as f64;

                total_norm_sq += grad_norm_sq;
            }
        }

        Ok(total_norm_sq.sqrt())
    }

    /// Compute losses WITHOUT updating network weights (for validation)
    ///
    /// This method is used during hyperparameter optimization to compute
    /// validation losses on held-out trajectories without updating the model.
    ///
    /// # Arguments
    /// * `batch` - Trajectory batch to compute losses on
    ///
    /// # Returns
    /// Tuple of (policy_loss, value_loss) as scalars
    pub fn compute_losses(&self, batch: &mut TrajectoryBatch) -> Result<(f32, f32), MLError> {
        // Convert batch to tensors
        let device = self.actor.device();
        let batch_tensors = batch.to_tensors(device, self.config.state_dim)?;

        // Compute losses WITHOUT backpropagation
        let policy_loss = self.compute_policy_loss(&batch_tensors)?;
        let value_loss = self.compute_value_loss(&batch_tensors)?;

        let policy_loss_scalar = policy_loss.to_scalar::<f32>()?;
        let value_loss_scalar = value_loss.to_scalar::<f32>()?;

        Ok((policy_loss_scalar, value_loss_scalar))
    }

    /// Compute `PPO` policy loss with clipping
    fn compute_policy_loss(&self, batch: &TrajectoryTensors) -> Result<Tensor, MLError> {
        // Get current log probabilities
        let new_log_probs = self.actor.log_probs(&batch.states, &batch.actions)?;

        // Compute probability ratio
        let log_ratio = (&new_log_probs - &batch.log_probs)?;
        let ratio = log_ratio.exp()?;

        // Clipped surrogate objective
        let clip_epsilon_tensor = Tensor::from_vec(
            vec![self.config.clip_epsilon; batch.advantages.dims()[0]],
            batch.advantages.dims(),
            self.actor.device(),
        )
        .map_err(|e| MLError::TrainingError(format!("Failed to create clip tensor: {}", e)))?;

        let one_tensor = Tensor::ones(batch.advantages.dims(), DType::F32, self.actor.device())?;
        let clip_min = (&one_tensor - &clip_epsilon_tensor)?;
        let clip_max = if let Some(eps_high) = self.config.clip_epsilon_high {
            let clip_high_tensor = Tensor::from_vec(
                vec![eps_high; batch.advantages.dims()[0]],
                batch.advantages.dims(),
                self.actor.device(),
            ).map_err(|e| MLError::TrainingError(format!("Failed to create clip_high tensor: {}", e)))?;
            (&one_tensor + &clip_high_tensor)?
        } else {
            (&one_tensor + &clip_epsilon_tensor)?
        };

        // Clamp ratio to [1-ε, 1+ε] (or [1-ε, 1+ε_high] with clip-higher)
        let clipped_ratio = ratio.clamp(&clip_min, &clip_max)?;

        // PPO objective: min(ratio * advantage, clipped_ratio * advantage)
        let surr1 = (&ratio * &batch.advantages)?;
        let surr2 = (&clipped_ratio * &batch.advantages)?;
        let policy_loss_raw = TensorOps::elementwise_min(&surr1, &surr2)?;

        // Add entropy bonus
        let entropy = self.actor.entropy(&batch.states)?;
        let entropy_bonus = TensorOps::scalar_mul(&entropy, self.config.entropy_coeff as f64)?;

        // Final loss (negative because we want to maximize)
        let policy_loss_inner = (policy_loss_raw + entropy_bonus)?.mean_all()?;
        let policy_loss = TensorOps::negate(&policy_loss_inner)?;

        Ok(policy_loss)
    }

    /// Compute value function loss
    fn compute_value_loss(&self, batch: &TrajectoryTensors) -> Result<Tensor, MLError> {
        let predicted_values = self.critic.forward(&batch.states)?;
        let value_loss = (&predicted_values - &batch.returns)?
            .powf(2.0)?
            .mean_all()?;
        let scaled_loss = TensorOps::scalar_mul(&value_loss, self.config.value_loss_coeff as f64)?;

        Ok(scaled_loss)
    }

    /// Initialize optimizers
    fn init_optimizers(&mut self) -> Result<(), MLError> {
        if self.policy_optimizer.is_none() {
            let policy_params = ParamsAdam {
                lr: self.config.policy_learning_rate,
                beta_1: 0.9,
                beta_2: 0.999,
                eps: 1e-8,
                weight_decay: None,
                amsgrad: false,
            };
            self.policy_optimizer = Some(
                Adam::new(self.actor.vars().all_vars(), policy_params).map_err(|e| {
                    MLError::TrainingError(format!("Failed to create policy optimizer: {}", e))
                })?,
            );
        }

        if self.value_optimizer.is_none() {
            let value_params = ParamsAdam {
                lr: self.config.value_learning_rate,
                beta_1: 0.9,
                beta_2: 0.999,
                eps: 1e-8,
                weight_decay: None,
                amsgrad: false,
            };
            self.value_optimizer = Some(
                Adam::new(self.critic.vars().all_vars(), value_params).map_err(|e| {
                    MLError::TrainingError(format!("Failed to create value optimizer: {}", e))
                })?,
            );
        }

        Ok(())
    }

    /// Get training steps
    pub fn get_training_steps(&self) -> u64 {
        self.training_steps
    }

    /// Get configuration
    pub fn get_config(&self) -> &PPOConfig {
        &self.config
    }

    /// Save PPO checkpoint with metadata (including training step counter)
    ///
    /// # Arguments
    /// * `actor_path` - Path to save actor (policy) network safetensors file
    /// * `critic_path` - Path to save critic (value) network safetensors file
    /// * `metadata_path` - Path to save metadata JSON file
    ///
    /// # Metadata Format
    /// ```json
    /// {
    ///   "training_steps": 1234,
    ///   "config": {
    ///     "state_dim": 64,
    ///     "num_actions": 3,
    ///     ...
    ///   }
    /// }
    /// ```
    pub fn save_checkpoint(
        &self,
        actor_path: &str,
        critic_path: &str,
        metadata_path: &str,
    ) -> Result<(), MLError> {
        // Save actor weights
        self.actor
            .vars()
            .save(actor_path)
            .map_err(|e| MLError::ModelError(format!("Failed to save actor checkpoint: {}", e)))?;

        // Save critic weights
        self.critic
            .vars()
            .save(critic_path)
            .map_err(|e| MLError::ModelError(format!("Failed to save critic checkpoint: {}", e)))?;

        // Save metadata with training_steps
        let metadata = serde_json::json!({
            "training_steps": self.training_steps,
            "config": {
                "state_dim": self.config.state_dim,
                "num_actions": self.config.num_actions,
                "policy_hidden_dims": self.config.policy_hidden_dims,
                "value_hidden_dims": self.config.value_hidden_dims,
                "policy_learning_rate": self.config.policy_learning_rate,
                "value_learning_rate": self.config.value_learning_rate,
                "clip_epsilon": self.config.clip_epsilon,
                "value_loss_coeff": self.config.value_loss_coeff,
                "entropy_coeff": self.config.entropy_coeff,
                "batch_size": self.config.batch_size,
                "mini_batch_size": self.config.mini_batch_size,
                "num_epochs": self.config.num_epochs,
                "max_grad_norm": self.config.max_grad_norm,
            }
        });

        std::fs::write(
            metadata_path,
            serde_json::to_string_pretty(&metadata)
                .map_err(|e| MLError::ModelError(format!("Failed to serialize metadata: {}", e)))?,
        )
        .map_err(|e| MLError::ModelError(format!("Failed to write metadata file: {}", e)))?;

        info!(
            "PPO checkpoint saved: actor={}, critic={}, metadata={}, training_steps={}",
            actor_path, critic_path, metadata_path, self.training_steps
        );

        Ok(())
    }

    /// Load PPO model from checkpoints (actor and critic networks)
    ///
    /// # Arguments
    /// * `actor_checkpoint_path` - Path to actor (policy) network safetensors file
    /// * `critic_checkpoint_path` - Path to critic (value) network safetensors file
    /// * `config` - PPO configuration (must match training config)
    /// * `device` - Device to load model on (CPU or CUDA)
    ///
    /// # Returns
    /// PPO instance with loaded weights from checkpoints
    ///
    /// # Example
    /// ```no_run
    /// use foxhunt_ml::ppo::{PPO, PPOConfig};
    /// use candle_core::Device;
    ///
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let config = PPOConfig::default();
    /// let device = Device::Cpu;
    /// let ppo = PPO::load_checkpoint(
    ///     "checkpoints/ppo_actor_epoch_100.safetensors",
    ///     "checkpoints/ppo_critic_epoch_100.safetensors",
    ///     config,
    ///     device,
    /// )?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn load_checkpoint(
        actor_checkpoint_path: &str,
        critic_checkpoint_path: &str,
        config: PPOConfig,
        device: Device,
    ) -> Result<Self, MLError> {
        info!(
            "Loading PPO checkpoint from actor={}, critic={}",
            actor_checkpoint_path, critic_checkpoint_path
        );

        // Extract values from config before moving it
        let use_lstm = config.use_lstm;
        let lstm_hidden_dim = config.lstm_hidden_dim;
        let lstm_num_layers = config.lstm_num_layers;

        // Load actor network from safetensors
        let actor_path = PathBuf::from(actor_checkpoint_path);
        // SAFETY: VarBuilder::from_mmaped_safetensors is safe here because:
        // 1. File path comes from user input and is validated by the safetensors deserializer
        // 2. Safetensors format guarantees correct memory layout (self-describing binary format)
        // 3. DType::F32 matches our checkpoint format (enforced during save)
        // 4. Memory-mapped access is read-only; file won't be modified during load
        // 5. Candle's SafeTensors deserializer validates the file format before creating tensors
        // 6. Any format violations cause an Err return, not undefined behavior
        //
        // The unsafe is inherited from memmap2::MmapOptions and is necessary for:
        // - Zero-copy deserialization (critical for HFT performance)
        // - Large model support (checkpoint files can be 100MB+)
        // - Avoiding full file read into memory
        //
        // Alternative: VarBuilder::from_buffered_safetensors loads into memory (safe but slower)
        //
        // SAFETY: Memory-mapped file access is safe because:
        // 1. File has just been verified to exist and is readable
        // 2. SafeTensors format includes checksums and validation
        // 3. Memory-mapped access is read-only (no modifications)
        // 4. Candle's deserializer validates format before tensor creation
        // 5. Any format violations cause Err return, not UB
        let actor_vb = unsafe {
            VarBuilder::from_mmaped_safetensors(&[actor_path], DType::F32, &device).map_err(
                |e| {
                    MLError::ModelError(format!(
                        "Failed to load actor checkpoint from {}: {}",
                        actor_checkpoint_path, e
                    ))
                },
            )?
        };

        let actor = if use_lstm {
            info!(
                "Loading LSTM actor checkpoint (hidden_dim={}, num_layers={})",
                lstm_hidden_dim, lstm_num_layers
            );
            let lstm_actor = LSTMPolicyNetwork::from_varbuilder(&config, actor_vb, &device)
                .map_err(|e| {
                    MLError::ModelError(format!(
                        "Failed to load LSTM actor from checkpoint: {}", e
                    ))
                })?;
            ActorNetwork::LSTM(lstm_actor)
        } else {
            let mlp_actor = PolicyNetwork::from_varbuilder(
                actor_vb,
                config.state_dim,
                &config.policy_hidden_dims,
                config.num_actions,
                device.clone(),
            )?;
            ActorNetwork::MLP(mlp_actor)
        };

        // Load critic network from safetensors
        let critic_path = PathBuf::from(critic_checkpoint_path);
        // SAFETY: VarBuilder::from_mmaped_safetensors is safe here because:
        // 1. File path comes from user input and is validated by the safetensors deserializer
        // 2. Safetensors format guarantees correct memory layout (self-describing binary format)
        // 3. DType::F32 matches our checkpoint format (enforced during save)
        // 4. Memory-mapped access is read-only; file won't be modified during load
        // 5. Candle's SafeTensors deserializer validates the file format before creating tensors
        // 6. Any format violations cause an Err return, not undefined behavior
        //
        // The unsafe is inherited from memmap2::MmapOptions and is necessary for:
        // - Zero-copy deserialization (critical for HFT performance)
        // - Large model support (checkpoint files can be 100MB+)
        // - Avoiding full file read into memory
        //
        // Alternative: VarBuilder::from_buffered_safetensors loads into memory (safe but slower)
        //
        // SAFETY: Memory-mapped file access is safe because:
        // 1. File has just been verified to exist and is readable
        // 2. SafeTensors format includes checksums and validation
        // 3. Memory-mapped access is read-only (no modifications)
        // 4. Candle's deserializer validates format before tensor creation
        // 5. Any format violations cause Err return, not UB
        let critic_vb = unsafe {
            VarBuilder::from_mmaped_safetensors(&[critic_path], DType::F32, &device).map_err(
                |e| {
                    MLError::ModelError(format!(
                        "Failed to load critic checkpoint from {}: {}",
                        critic_checkpoint_path, e
                    ))
                },
            )?
        };

        let critic = if use_lstm {
            info!(
                "Loading LSTM critic checkpoint (hidden_dim={}, num_layers={})",
                lstm_hidden_dim, lstm_num_layers
            );
            let lstm_critic = LSTMValueNetwork::from_varbuilder(&config, critic_vb, &device)
                .map_err(|e| {
                    MLError::ModelError(format!(
                        "Failed to load LSTM critic from checkpoint: {}", e
                    ))
                })?;
            CriticNetwork::LSTM(lstm_critic)
        } else {
            let mlp_critic = ValueNetwork::from_varbuilder(
                critic_vb,
                config.state_dim,
                &config.value_hidden_dims,
                device.clone(),
            )?;
            CriticNetwork::MLP(mlp_critic)
        };

        // Try to load metadata to restore training_steps
        // Look for metadata file in same directory as actor checkpoint
        let actor_path = std::path::Path::new(actor_checkpoint_path);
        let metadata_path = if let Some(parent) = actor_path.parent() {
            if let Some(stem) = actor_path.file_stem() {
                // Try metadata file matching actor checkpoint name pattern
                parent.join(format!("{}_metadata.json", stem.to_string_lossy()))
            } else {
                parent.join("checkpoint_metadata.json")
            }
        } else {
            PathBuf::from("checkpoint_metadata.json")
        };

        let training_steps = if metadata_path.exists() {
            match std::fs::read_to_string(&metadata_path) {
                Ok(metadata_str) => {
                    match serde_json::from_str::<serde_json::Value>(&metadata_str) {
                        Ok(metadata) => {
                            let steps = metadata
                                .get("training_steps")
                                .and_then(|v| v.as_u64())
                                .unwrap_or(0);
                            info!(
                                "Restored training_steps={} from metadata file: {:?}",
                                steps, metadata_path
                            );
                            steps
                        },
                        Err(e) => {
                            warn!(
                            "Failed to parse metadata JSON from {:?}: {}. Starting from step 0.",
                            metadata_path, e
                        );
                            0
                        },
                    }
                },
                Err(e) => {
                    warn!(
                        "Failed to read metadata file {:?}: {}. Starting from step 0.",
                        metadata_path, e
                    );
                    0
                },
            }
        } else {
            info!(
                "No metadata file found at {:?}. Starting from step 0 (legacy checkpoint).",
                metadata_path
            );
            0
        };

        info!(
            "PPO checkpoint loaded successfully with training_steps={}",
            training_steps
        );

        // Initialize risk management components
        let portfolio_tracker = PortfolioTracker::new(10_000.0, 0.0001, config.cash_reserve_pct);
        let reward_normalizer = Some(RewardNormalizer::new());
        let circuit_breaker_config = CircuitBreakerConfig {
            failure_threshold: config.circuit_breaker_threshold,
            success_threshold: 3,
            timeout_duration: std::time::Duration::from_secs(60),
            half_open_max_calls: 2,
        };
        let circuit_breaker = Some(CircuitBreaker::new(circuit_breaker_config));

        // Extract values before moving config
        let transaction_cost_bps = config.transaction_cost_bps;
        let max_position_absolute = config.max_position_absolute;
        // Note: use_lstm, lstm_hidden_dim, lstm_num_layers already extracted at function start

        // Initialize hidden state manager if LSTM is enabled
        // LSTM hidden/cell states are not saved in checkpoints -- only network weights.
        // States are re-initialized to zeros on load via HiddenStateManager::new().
        let hidden_state_manager = use_lstm
            .then(|| {
                HiddenStateManager::new(
                    lstm_num_layers,
                    1, // Default batch size
                    lstm_hidden_dim,
                    &device,
                )
            })
            .transpose()?;

        Ok(Self {
            config,
            actor,  // Already wrapped in ActorNetwork enum variant
            critic,  // Already wrapped in CriticNetwork enum variant
            policy_optimizer: None,
            value_optimizer: None,
            training_steps,
            portfolio_tracker,
            reward_normalizer,
            circuit_breaker,
            transaction_cost_bps: Some(transaction_cost_bps),
            max_position_absolute: Some(max_position_absolute),
            hidden_state_manager,
        })
    }

    /// Predict action probabilities for a given state
    ///
    /// # Arguments
    /// * `state` - State vector (must match config.state_dim)
    ///
    /// # Returns
    /// Vector of action probabilities (length = config.num_actions)
    pub fn predict(&self, state: &[f32]) -> Result<Vec<f32>, MLError> {
        if state.len() != self.config.state_dim {
            return Err(MLError::InvalidInput(format!(
                "State dimension mismatch: expected {}, got {}",
                self.config.state_dim,
                state.len()
            )));
        }

        let state_tensor = Tensor::from_vec(
            state.to_vec(),
            (1, self.config.state_dim),
            self.actor.device(),
        )?;
        let probs_tensor = self.actor.action_probabilities(&state_tensor)?;
        let probs = probs_tensor.flatten_all()?.to_vec1::<f32>()?;
        Ok(probs)
    }

    /// Select the greedy (deterministic) action -- argmax of policy probabilities.
    ///
    /// Use this for validation/evaluation where stochastic sampling would
    /// introduce noise into the performance metric.
    pub fn greedy_action(&self, state: &[f32]) -> Result<TradingAction, MLError> {
        let probs = self.predict(state)?;
        let best_idx = probs
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, _)| i)
            .unwrap_or(2); // default HOLD
        TradingAction::from_int(best_idx as u8)
            .ok_or_else(|| MLError::InvalidInput(format!("Invalid action index: {}", best_idx)))
    }
}

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

    #[test]
    fn test_policy_network_creation() -> Result<()> {
        let device = Device::Cpu;
        let _policy = PolicyNetwork::new(10, &[32, 16], 3, device)
            .map_err(|_| anyhow::anyhow!("Failed to create policy network"))?;
        // Policy network created successfully
        Ok(())
    }

    #[test]
    fn test_value_network_creation() -> Result<()> {
        let device = Device::Cpu;
        let _value = ValueNetwork::new(10, &[32, 16], device)
            .map_err(|_| anyhow::anyhow!("Failed to create value network"))?;
        // Value network created successfully
        Ok(())
    }

    #[test]
    fn test_ppo_creation() -> Result<()> {
        let config = PPOConfig::default();
        let ppo = PPO::new(config).map_err(|_| anyhow::anyhow!("Failed to create PPO"))?;
        // PPO created successfully
        assert_eq!(ppo.get_training_steps(), 0);
        Ok(())
    }

    #[test]
    fn test_ppo_config_default() -> Result<()> {
        let config = PPOConfig::default();
        assert!(config.state_dim > 0);
        assert!(config.num_actions > 0);
        assert!(config.policy_learning_rate > 0.0);
        assert!(config.value_learning_rate > 0.0);
        Ok(())
    }

    #[test]
    fn test_ppo_training_steps() -> Result<()> {
        let config = PPOConfig::default();
        let mut ppo =
            PPO::new(config).map_err(|_| anyhow::anyhow!("Failed to create PPO"))?;

        assert_eq!(ppo.get_training_steps(), 0);
        ppo.training_steps = 5;
        assert_eq!(ppo.get_training_steps(), 5);
        Ok(())
    }

    #[test]
    fn test_ppo_config_validation() -> Result<()> {
        let config = PPOConfig {
            clip_epsilon: 0.2,
            value_loss_coeff: 1.0,
            entropy_coeff: 0.05,
            ..Default::default()
        };

        let ppo = PPO::new(config).map_err(|_| anyhow::anyhow!("Failed to create PPO"))?;
        assert_eq!(ppo.get_config().clip_epsilon, 0.2);
        assert_eq!(ppo.get_config().value_loss_coeff, 1.0);
        Ok(())
    }
}