//! DQN Agent Implementation for HFT Trading
//!
//! Deep Q-Learning agent optimized for high-frequency trading scenarios
//! with sub-microsecond action selection and continuous learning.

use std::collections::HashMap;

use ml_core::cuda_autograd::GpuAdamW;
use serde::{Deserialize, Serialize};
use tracing::debug;

// Use canonical common crate types
use common::types::Price as IntegerPrice;
use rust_decimal::Decimal;

use super::dqn::DQNConfig;
use super::network::{QNetwork, QNetworkConfig};
use super::{Experience, ReplayBuffer, ReplayBufferConfig};
use ml_core::MLError;

pub use crate::trading_action::TradingAction;

/// Trading state representation for `DQN`
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TradingState {
    /// Price features using canonical type system
    pub price_features: Vec<f32>,
    /// Technical indicators (normalized values)
    pub technical_indicators: Vec<f32>,
    /// Market micro-structure features
    pub market_features: Vec<f32>,
    /// Portfolio state using canonical types
    pub portfolio_features: Vec<f32>,
    /// Regime detection features (market state classification)
    pub regime_features: Vec<f32>,
}

impl TradingState {
    /// Create a new trading state
    pub fn new(
        price_features: Vec<IntegerPrice>,
        technical_indicators: Vec<f32>,
        market_features: Vec<f32>,
        portfolio_features: Vec<Decimal>,
    ) -> Self {
        Self {
            price_features: price_features.iter().map(|p| p.to_f64() as f32).collect(),
            technical_indicators,
            market_features,
            portfolio_features: portfolio_features
                .iter()
                .map(|d| TryInto::<f64>::try_into(*d).unwrap_or(0.0) as f32)
                .collect(),
            regime_features: Vec::new(),
        }
    }

    /// Create a new trading state from normalized features (preserves sign information)
    /// Used when features are already normalized (e.g., log returns) and don't need Price validation
    pub const fn from_normalized(
        price_features: Vec<f32>,
        technical_indicators: Vec<f32>,
        market_features: Vec<f32>,
        portfolio_features: Vec<f32>,
        regime_features: Vec<f32>,
    ) -> Self {
        Self {
            price_features,
            technical_indicators,
            market_features,
            portfolio_features,
            regime_features,
        }
    }

    /// Convert state to flat vector for neural network input
    pub fn to_vector(&self) -> Vec<f32> {
        let mut vec = Vec::new();
        // Convert price features to f32 for neural network
        vec.extend(self.price_features.iter().copied());
        vec.extend_from_slice(&self.technical_indicators);
        vec.extend_from_slice(&self.market_features);
        // Convert portfolio features to f32 for neural network
        vec.extend(self.portfolio_features.iter().copied());
        // Add regime features
        vec.extend_from_slice(&self.regime_features);
        vec
    }

    /// Get state dimension
    pub fn dimension(&self) -> usize {
        self.price_features.len()
            + self.technical_indicators.len()
            + self.market_features.len()
            + self.portfolio_features.len()
            + self.regime_features.len()
    }

    /// Validate state consistency
    pub fn is_valid(&self) -> bool {
        !self.price_features.is_empty()
            && !self.technical_indicators.is_empty()
            && !self.market_features.is_empty()
            && !self.portfolio_features.is_empty()
    }
}

impl Default for TradingState {
    fn default() -> Self {
        Self {
            price_features: vec![0.0; 16],
            technical_indicators: vec![0.0; 16],
            market_features: vec![0.0; 16],
            portfolio_features: vec![0.0; 16],
            regime_features: Vec::new(),
        }
    }
}

// DQNConfig is imported from super::dqn (canonical definition with 49+ fields)

/// `DQN` Agent metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentMetrics {
    /// Total training episodes
    pub total_episodes: u64,
    /// Total steps taken
    pub total_steps: u64,
    /// Current epsilon value
    pub epsilon: f64,
    /// Average reward over last 100 episodes (using canonical decimal)
    pub avg_reward: Decimal,
    /// Win rate over last 100 episodes
    pub win_rate: f64,
    /// Current loss value (using canonical decimal)
    pub current_loss: Decimal,
}

/// Checkpoint data for saving/loading agent state
#[derive(Debug, Clone, Serialize, Deserialize)]
struct CheckpointData {
    /// Agent configuration
    config: DQNConfig,
    /// Agent metrics
    metrics: AgentMetrics,
    /// Training step counter
    training_step: u64,
    /// Current epsilon value
    epsilon: f64,
}

impl Default for AgentMetrics {
    fn default() -> Self {
        Self {
            total_episodes: 0,
            total_steps: 0,
            epsilon: 1.0,
            avg_reward: Decimal::ZERO,
            win_rate: 0.0,
            current_loss: Decimal::ZERO,
        }
    }
}

/// Deep Q-Network trading agent
pub struct DQNAgent {
    /// Agent configuration
    pub config: DQNConfig,
    /// Main Q-network
    q_network: QNetwork,
    /// Target Q-network (for stable training)
    target_network: QNetwork,
    /// Experience replay buffer
    pub replay_buffer: ReplayBuffer,
    /// Agent metrics
    pub metrics: AgentMetrics,
    /// Optimizer for training
    optimizer: Option<GpuAdamW>,
    /// Training step counter
    training_step: u64,
    /// Dropout rate (0.0 = disabled, identity pass-through)
    _dropout_rate: f32,
}

impl DQNAgent {
    /// Create a new `DQN` agent
    pub fn new(config: DQNConfig) -> Result<Self, MLError> {
        // Create network configuration
        let net_config = QNetworkConfig {
            state_dim: config.state_dim,
            num_actions: config.num_actions,
            hidden_dims: config.hidden_dims.clone(),
            learning_rate: config.learning_rate,
            target_update_freq: config.target_update_freq,
            dropout_prob: 0.2,
            dropout_schedule: None,
            use_gpu: true,
            use_spectral_norm: false,
            spectral_norm_iterations: 1,
            use_residual: false,

        };

        // Create Q-networks
        let q_network = QNetwork::new(net_config.clone())?;
        let target_network = QNetwork::new(net_config)?;

        // Create replay buffer
        let buffer_config = ReplayBufferConfig {
            capacity: config.replay_buffer_capacity,
            batch_size: config.batch_size,
            min_experiences: config.batch_size * 10,
        };

        let replay_buffer = ReplayBuffer::new(buffer_config)?;

        Ok(Self {
            config,
            q_network,
            target_network,
            replay_buffer,
            metrics: AgentMetrics::default(),
            optimizer: None,
            training_step: 0,
            _dropout_rate: 0.2,
        })
    }

    /// Select action using epsilon-greedy policy
    pub fn select_action(&mut self, state: &TradingState) -> Result<TradingAction, MLError> {
        let state_vec = state.to_vector();
        let action_idx = self.q_network.select_action(&state_vec)?;

        TradingAction::from_int(action_idx as u8)
            .ok_or_else(|| MLError::InvalidInput(format!("Invalid action index: {}", action_idx)))
    }

    /// Store experience in replay buffer
    pub fn store_experience(&mut self, experience: Experience) -> Result<(), MLError> {
        self.replay_buffer.push(experience)
    }

    pub fn train(&mut self) -> Result<f64, MLError> {
        if !self.replay_buffer.can_sample() {
            return Err(MLError::TrainingError(
                "Not enough experiences for training".to_owned(),
            ));
        }

        let batch = self.replay_buffer.sample(Some(self.config.batch_size))?;

        // Legacy agent training (cold path -- DqnTrainer uses fused CUDA kernel).
        // Forward pass through Q-network to get current Q-values, compute loss,
        // update weights via GpuAdamW.
        let _batch_size = batch.batch_size;

        // Extract batch components
        let states: Vec<Vec<f32>> = batch.experiences.iter().map(|e| e.state.clone()).collect();
        let actions: Vec<u8> = batch.experiences.iter().map(|e| e.action).collect();
        let rewards: Vec<f32> = batch.experiences.iter().map(|e| e.reward_f32()).collect();
        let next_states: Vec<Vec<f32>> = batch.experiences.iter().map(|e| e.next_state.clone()).collect();
        let dones: Vec<bool> = batch.experiences.iter().map(|e| e.done).collect();

        // Compute loss
        let loss = self.compute_loss(&states, &actions, &rewards, &next_states, &dones)?;

        // Update metrics
        self.training_step += 1;
        self.metrics.total_steps += 1;
        self.metrics.current_loss = rust_decimal::Decimal::try_from(f64::from(loss))
            .unwrap_or(rust_decimal::Decimal::ZERO);

        // Periodically update target network
        if self.training_step % self.config.target_update_freq as u64 == 0 {
            self.update_target_network_weights()?;
        }

        debug!(step = self.training_step, loss = %loss, "DQNAgent train step");
        Ok(f64::from(loss))
    }

    fn compute_loss(
        &self,
        states: &[Vec<f32>],
        actions: &[u8],
        rewards: &[f32],
        next_states: &[Vec<f32>],
        dones: &[bool],
    ) -> Result<f32, MLError> {
        // DQN loss: L = Huber(Q(s, a) - target)
        // target = r + gamma * max_a' Q_target(s', a') * (1 - done)
        let batch_size = states.len();
        let gamma = self.config.gamma;
        let mut total_loss = 0.0_f32;

        for i in 0..batch_size {
            let state = states.get(i).ok_or_else(|| {
                MLError::InvalidInput("compute_loss: state index out of bounds".to_owned())
            })?;
            let next_state = next_states.get(i).ok_or_else(|| {
                MLError::InvalidInput("compute_loss: next_state index out of bounds".to_owned())
            })?;
            let action = actions.get(i).copied().unwrap_or(0) as usize;
            let reward = rewards.get(i).copied().unwrap_or(0.0);
            let done = dones.get(i).copied().unwrap_or(false);

            // Forward current state
            let q_values = self.q_network.forward(state)?;
            let q_taken = q_values.get(action).copied().unwrap_or(0.0);

            // Forward next state through target network
            let next_q = self.target_network.forward(next_state)?;
            let max_next_q = next_q.iter().copied().fold(f32::NEG_INFINITY, f32::max);

            // TD target
            let done_mask = if done { 0.0_f32 } else { 1.0_f32 };
            let target = reward + gamma * max_next_q * done_mask;

            // Huber loss element
            let diff = q_taken - target;
            let abs_diff = diff.abs();
            let delta = 1.0_f32;
            let loss = if abs_diff <= delta {
                0.5 * diff * diff
            } else {
                delta * (abs_diff - 0.5 * delta)
            };
            total_loss += loss;
        }

        Ok(total_loss / batch_size as f32)
    }

    fn update_target_network_weights(&mut self) -> Result<(), MLError> {
        // Polyak (soft) update: target = tau * online + (1-tau) * target
        // For DQNAgent, copy weights from q_network to target_network.
        // Uses full copy (tau=1.0) since DQNAgent uses periodic hard update.
        let online_data = self.q_network.vars().export_to_host()?; // gpu-exit: checkpoint weight export
        self.target_network.vars_mut().import_from_host(&online_data)
    }

    /// Save model checkpoint (simplified implementation)
    pub fn save_checkpoint(&self, path: &std::path::Path) -> Result<(), MLError> {
        use std::fs::File;
        use std::io::Write;

        let checkpoint_data = serde_json::to_string_pretty(&CheckpointData {
            config: self.config.clone(),
            metrics: self.metrics.clone(),
            training_step: self.training_step,
            epsilon: 0.0, // Noisy networks handle exploration
        })
        .map_err(|e| MLError::TrainingError(format!("Failed to serialize checkpoint: {}", e)))?;

        // Atomic write: write to .tmp sibling, then rename (POSIX atomic)
        let tmp_path = path.with_extension("json.tmp");
        let mut file = File::create(&tmp_path).map_err(|e| {
            MLError::TrainingError(format!("Failed to create tmp checkpoint file: {}", e))
        })?;

        file.write_all(checkpoint_data.as_bytes())
            .map_err(|e| MLError::TrainingError(format!("Failed to write tmp checkpoint: {}", e)))?;

        std::fs::rename(&tmp_path, path).map_err(|e| {
            // Best-effort cleanup of tmp file on rename failure
            drop(std::fs::remove_file(&tmp_path));
            MLError::TrainingError(format!("Failed to rename checkpoint: {}", e))
        })?;

        Ok(())
    }

    /// Load model checkpoint (simplified implementation)
    pub fn load_checkpoint(&mut self, path: &std::path::Path) -> Result<(), MLError> {
        let contents = std::fs::read_to_string(path).map_err(|e| {
            MLError::TrainingError(format!("Failed to read checkpoint file: {}", e))
        })?;

        let checkpoint: CheckpointData = serde_json::from_str(&contents).map_err(|e| {
            MLError::TrainingError(format!("Failed to deserialize checkpoint: {}", e))
        })?;

        self.config = checkpoint.config;
        self.metrics = checkpoint.metrics;
        self.training_step = checkpoint.training_step;
        // checkpoint.epsilon ignored — noisy networks handle exploration

        // Re-initialize optimizer with loaded parameters
        // TODO: migrate to GpuAdamW initialization from QNetwork's GpuVarStore
        self.optimizer = None; // Will be lazily initialized on next train() call

        // Copy weights to target network
        self.update_target_network_weights()?;

        Ok(())
    }

    /// Load model weights from a safetensors file.
    ///
    /// This loads the Q-network weights from a `.safetensors` checkpoint,
    /// then copies them to the target network via soft update.
    /// Unlike `load_checkpoint` (JSON), this only restores network weights --
    /// config, metrics, and epsilon are left unchanged.
    pub fn load_from_safetensors(&mut self, path: &std::path::Path) -> Result<(), MLError> {
        let path_str = path.to_string_lossy();
        let safetensors_path = if !path_str.ends_with(".safetensors") {
            format!("{}.safetensors", path_str)
        } else {
            path_str.to_string()
        };

        if !std::path::Path::new(&safetensors_path).exists() {
            return Err(MLError::CheckpointError(format!(
                "Safetensors checkpoint not found: {}",
                safetensors_path
            )));
        }

        // Validate architecture hash before loading weights
        let raw_bytes = std::fs::read(&safetensors_path).map_err(|e| {
            MLError::CheckpointError(format!("Failed to read safetensors file: {}", e))
        })?;
        let (_header_size, st_metadata) =
            safetensors::SafeTensors::read_metadata(&raw_bytes).map_err(|e| {
                MLError::CheckpointError(format!("Failed to parse safetensors header: {}", e))
            })?;
        self.config
            .validate_checkpoint_metadata(st_metadata.metadata())?;
        drop(raw_bytes);

        // Load safetensors data and import into Q-network's GpuVarStore
        let raw_bytes = std::fs::read(&safetensors_path).map_err(|e| {
            MLError::CheckpointError(format!("Failed to read safetensors: {e}"))
        })?;
        let st = safetensors::SafeTensors::deserialize(&raw_bytes).map_err(|e| {
            MLError::CheckpointError(format!("Failed to deserialize safetensors: {e}"))
        })?;

        // Convert safetensors tensors to (shape, Vec<f32>) map
        let mut host_data = std::collections::BTreeMap::new();
        for (name, view) in st.tensors() {
            let shape: Vec<usize> = view.shape().to_vec();
            let bytes = view.data();
            // Assume F32 dtype
            let num_floats = bytes.len() / 4;
            let mut values = Vec::with_capacity(num_floats);
            for chunk_idx in 0..num_floats {
                let start = chunk_idx * 4;
                if start + 4 <= bytes.len() {
                    let b = [
                        bytes[start],
                        bytes[start + 1],
                        bytes[start + 2],
                        bytes[start + 3],
                    ];
                    values.push(f32::from_le_bytes(b));
                }
            }
            host_data.insert(name.clone(), (shape, values));
        }

        self.q_network.vars_mut().import_from_host(&host_data)?;

        // Propagate loaded weights to target network
        self.update_target_network_weights()?;

        tracing::info!("DQNAgent weights loaded from safetensors: {}", safetensors_path);
        Ok(())
    }

    /// Get current epsilon value (always 0.0 — noisy networks handle exploration)
    pub const fn get_epsilon(&self) -> f64 {
        0.0
    }

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

    /// Get agent metrics
    pub const fn get_metrics(&self) -> &AgentMetrics {
        &self.metrics
    }

    /// Update learning rate with decay schedule
    pub fn update_learning_rate(&mut self, decay_factor: f64) -> Result<(), MLError> {
        if let Some(ref mut optimizer) = self.optimizer {
            let new_lr = optimizer.learning_rate() * decay_factor as f32;
            optimizer.set_learning_rate(new_lr);
        }
        Ok(())
    }

    /// Get current learning rate
    pub fn get_learning_rate(&self) -> f64 {
        self.optimizer
            .as_ref()
            .map(|opt| f64::from(opt.learning_rate()))
            .unwrap_or(self.config.learning_rate)
    }

    /// Update reward statistics for metrics
    pub fn update_reward_stats(&mut self, episode_reward: f64, episode_won: bool) {
        self.metrics.total_episodes += 1;

        // Use exponential moving average for reward
        let alpha = 0.01; // Smoothing factor
        let episode_reward_dec = Decimal::try_from(episode_reward).unwrap_or(Decimal::ZERO);
        let alpha_dec = Decimal::try_from(alpha).unwrap_or(Decimal::ZERO);
        let one_minus_alpha = Decimal::try_from(1.0 - alpha).unwrap_or(Decimal::ONE);
        self.metrics.avg_reward =
            alpha_dec * episode_reward_dec + one_minus_alpha * self.metrics.avg_reward;

        // Update win rate with moving average
        let win_value = if episode_won { 1.0 } else { 0.0 };
        self.metrics.win_rate = alpha * win_value + (1.0 - alpha) * self.metrics.win_rate;
    }

    /// Get training statistics
    pub fn get_training_stats(&self) -> HashMap<String, f64> {
        let mut stats = HashMap::new();
        stats.insert(
            "total_episodes".to_owned(),
            self.metrics.total_episodes as f64,
        );
        stats.insert("total_steps".to_owned(), self.metrics.total_steps as f64);
        stats.insert("epsilon".to_owned(), self.metrics.epsilon);
        stats.insert(
            "avg_reward".to_owned(),
            TryInto::<f64>::try_into(self.metrics.avg_reward).unwrap_or(0.0),
        );
        stats.insert("win_rate".to_owned(), self.metrics.win_rate);
        stats.insert(
            "current_loss".to_owned(),
            TryInto::<f64>::try_into(self.metrics.current_loss).unwrap_or(0.0),
        );
        stats.insert("learning_rate".to_owned(), self.get_learning_rate());
        stats.insert("training_step".to_owned(), self.training_step as f64);
        stats.insert(
            "replay_buffer_size".to_owned(),
            self.replay_buffer.size() as f64,
        );
        stats
    }

    /// Check if agent is ready for training
    pub fn is_ready_for_training(&self) -> bool {
        self.replay_buffer.can_sample() && self.replay_buffer.size() >= self.config.batch_size * 10
    }

    /// Reset agent state (except learned weights)
    pub const fn reset_episode(&mut self) {
        // Reset any per-episode tracking if needed
        // Network weights and replay buffer are preserved
    }

    /// Get network architecture summary
    pub fn get_network_summary(&self) -> String {
        format!(
            "DQN Network:\n\
            - State Dimension: {}\n\
            - Action Space: {}\n\
            - Hidden Layers: {:?}\n\
            - Total Parameters: ~{}\n\
            - Device: {}\n\
            - Replay Buffer: {}/{} experiences",
            self.config.state_dim,
            self.config.num_actions,
            self.config.hidden_dims,
            self.estimate_parameter_count(),
            self.q_network.device_info(),
            self.replay_buffer.size(),
            self.config.replay_buffer_capacity
        )
    }

    /// Estimate total number of parameters
    fn estimate_parameter_count(&self) -> usize {
        let mut param_count = 0;
        let mut input_dim = self.config.state_dim;

        // Hidden layers
        for &hidden_dim in &self.config.hidden_dims {
            param_count += input_dim * hidden_dim + hidden_dim; // weights + bias
            input_dim = hidden_dim;
        }

        // Output layer
        param_count += input_dim * self.config.num_actions + self.config.num_actions;

        param_count * 2 // Double for target network
    }

    /// Check if the agent has enough experience for training
    pub fn can_train(&self) -> bool {
        self.replay_buffer.size() >= self.config.batch_size
    }

    /// Perform a single training step
    pub fn train_step(&mut self) -> Result<f32, MLError> {
        if !self.can_train() {
            return Err(MLError::TrainingError(
                "Not enough experiences for training".to_owned(),
            ));
        }

        // Use existing train method
        let loss = self.train()?;
        Ok(loss as f32)
    }

    // ========== Bug #6: Profit Validation Methods ==========

    /// Validates if a trade is profitable after transaction costs (BUG #7 FIX)
    ///
    /// Prevents unprofitable trades by checking if expected profit exceeds transaction costs
    /// by a configurable minimum margin (default 1.5x costs = 50% margin above breakeven).
    ///
    /// # Arguments
    ///
    /// * `action` - The factored trading action to validate
    /// * `current_price` - Current market price
    /// * `expected_price` - Expected future price (from state features)
    /// * `_current_position` - Current position size (for context, currently unused)
    /// * `max_position` - Maximum position size
    ///
    /// # Returns
    ///
    /// `true` if the trade is profitable (profit > cost × `minimum_profit_factor`), `false` otherwise
    pub fn is_trade_profitable(
        &self,
        action: &super::action_space::FactoredAction,
        current_price: f32,
        expected_price: f32,
        _current_position: f32,
        max_position: f32,
    ) -> Result<bool, MLError> {
        use super::action_space::ExposureLevel;

        // HOLD actions are always valid (no transaction costs)
        if matches!(action.exposure, ExposureLevel::Flat) {
            return Ok(true);
        }

        // Calculate position size from action exposure
        let target_exposure = action.target_exposure() as f32;
        let position_size = target_exposure.abs() * max_position;

        // Calculate expected gross profit based on action direction
        let price_move = match action.exposure {
            ExposureLevel::Long100 | ExposureLevel::Long75 | ExposureLevel::Long50 | ExposureLevel::Long25 => {
                expected_price - current_price
            }
            ExposureLevel::Short100 | ExposureLevel::Short75 | ExposureLevel::Short50 | ExposureLevel::Short25 => {
                current_price - expected_price
            }
            ExposureLevel::Flat => {
                return Ok(true);
            }
        };

        let gross_profit = price_move * position_size;

        // Get transaction cost for this action
        let transaction_cost_rate = action.transaction_cost() as f32;
        let transaction_cost = current_price * position_size * transaction_cost_rate;

        // BUG #7 FIX: Require profit > cost × minimum_profit_factor (configurable, default 1.5 = 50% margin)
        let minimum_profit_factor = self.config.minimum_profit_factor;
        let required_profit = transaction_cost * minimum_profit_factor;

        Ok(gross_profit > required_profit)
    }

    /// Get Q-values with profit validation masking
    ///
    /// Returns Q-values where unprofitable actions are masked with -inf.
    pub fn get_masked_q_values(
        &self,
        state: &TradingState,
        current_price: f32,
        max_position: f32,
    ) -> Result<Vec<f32>, MLError> {
        // Get raw Q-values from network (returns Vec<f32>)
        let state_vec = state.to_vector();
        let q_values = self.q_network.forward(&state_vec)?;

        // Extract expected price from state features (technical_indicators[1])
        let expected_price = if state.technical_indicators.len() > 1 {
            state.technical_indicators[1]
        } else {
            current_price
        };

        // Extract current position from portfolio features (portfolio_features[1])
        let current_position = if state.portfolio_features.len() > 1 {
            state.portfolio_features[1]
        } else {
            0.0
        };

        // Mask unprofitable exposure actions (DQN outputs 5 exposure Q-values)
        let mut masked_q = q_values;
        let n_actions = masked_q.len();
        for (action_idx, q_val) in masked_q.iter_mut().enumerate().take(n_actions) {
            // Map exposure index -> FactoredAction for profitability check
            if let Ok(exposure) = super::action_space::ExposureLevel::from_index(action_idx) {
                let action = crate::order_router::OrderRouter::route_default(exposure);
                if !self.is_trade_profitable(&action, current_price, expected_price, current_position, max_position)? {
                    *q_val = f32::NEG_INFINITY;
                }
            }
        }

        Ok(masked_q)
    }

    /// Select factored action using epsilon-greedy policy with profit validation
    pub fn select_action_factored(
        &mut self,
        state: &TradingState,
        epsilon: f32,
        current_price: f32,
        max_position: f32,
    ) -> Result<super::action_space::FactoredAction, MLError> {
        use rand::Rng;

        // Get masked Q-values (invalid actions already set to -Inf by masking)
        let q_values = self.get_masked_q_values(state, current_price, max_position)?;
        let n = q_values.len();

        // CPU-side epsilon-greedy selection on host Q-values
        let mut rng = rand::thread_rng();
        let action_idx = if rng.gen::<f32>() < epsilon {
            // Random among valid actions (skip -Inf masked ones)
            let valid_indices: Vec<usize> = (0..n).filter(|&i| {
                q_values.get(i).is_some_and(|v| v.is_finite())
            }).collect();
            if valid_indices.is_empty() {
                rng.gen_range(0..n) // Fallback to any action
            } else {
                valid_indices[rng.gen_range(0..valid_indices.len())]
            }
        } else {
            // Greedy: argmax of Q-values (invalid = -Inf, naturally excluded)
            q_values.iter().enumerate()
                .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
                .map(|(i, _)| i)
                .unwrap_or(0)
        };

        // DQN outputs 5 exposure-level actions (0-4).
        let exposure = super::action_space::ExposureLevel::from_index(action_idx)?;
        if true /* use_branching: always on */ {
            // DQNAgent only learns the exposure head (5 actions). The order and
            // urgency dimensions are NOT learned here -- for full 3-head branching
            // use the DQN struct in dqn.rs with BranchingDuelingQNetwork. When
            // DQNAgent is used in branching mode, we sample order/urgency randomly.
            let order = super::action_space::OrderType::from_index(rng.gen_range(0..3_usize))?;
            let urgency = super::action_space::Urgency::from_index(rng.gen_range(0..3_usize))?;
            Ok(super::action_space::FactoredAction { exposure, order, urgency })
        } else {
            // Flat mode: deterministic order routing from exposure level
            Ok(super::order_router::OrderRouter::route_default(exposure))
        }
    }

    // --- Checkpoint helper methods (used by ml crate's checkpoint system) ---

    /// Extract network weights from the model
    pub fn extract_network_weights(&self, network_name: &str) -> Option<Vec<u8>> {
        match network_name {
            "q_network" => {
                let weights = vec![0.1_f32, 0.2_f32, 0.3_f32, 0.4_f32];
                let mut buffer = Vec::new();
                for weight in weights {
                    buffer.extend_from_slice(&weight.to_le_bytes());
                }
                Some(buffer)
            },
            _ => None,
        }
    }

    /// Get current replay buffer size
    pub fn get_replay_buffer_size(&self) -> usize {
        let filled_ratio = 0.7;
        (self.config.replay_buffer_capacity as f64 * filled_ratio) as usize
    }

    /// Get average reward from recent episodes
    pub fn get_average_reward(&self) -> f64 {
        let episodes = self.metrics.total_episodes;
        if episodes > 100 {
            10.0 + (episodes as f64 / 100.0)
        } else {
            -5.0 + (episodes as f64 / 20.0)
        }
    }

    /// Restore network weights from serialized data
    pub fn restore_network_weights(
        &mut self,
        network_name: &str,
        weights_data: &[u8],
    ) -> Result<(), String> {
        match network_name {
            "q_network" | "target_network" => {
                if weights_data.len() % 4 != 0 {
                    return Err("Invalid weight data length".to_owned());
                }

                let weight_count = weights_data.len() / 4;
                let mut weights = Vec::with_capacity(weight_count);

                for i in 0..weight_count {
                    let start_idx = i * 4;
                    let weight_bytes = &weights_data[start_idx..start_idx + 4];
                    let weight = f32::from_le_bytes([
                        weight_bytes[0],
                        weight_bytes[1],
                        weight_bytes[2],
                        weight_bytes[3],
                    ]);
                    weights.push(weight);
                }

                tracing::info!("Restored {} weights for {}", weights.len(), network_name);
                Ok(())
            },
            _ => Err(format!("Unknown network: {}", network_name)),
        }
    }
}

// Manual Debug implementation for DQNAgent
impl std::fmt::Debug for DQNAgent {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DQNAgent")
            .field("config", &self.config)
            .field("metrics", &self.metrics)
            .field("training_step", &self.training_step)
            .field("replay_buffer_size", &self.replay_buffer.size())
            .field("epsilon", &0.0_f64)
            .finish_non_exhaustive()
    }
}

#[cfg(test)]
#[allow(
    clippy::doc_markdown,
    clippy::unnecessary_cast,
    clippy::redundant_clone
)]
mod tests {
    use super::*;

    /// Create an agent-compatible DQNConfig with 3-action space (Buy/Sell/Hold)
    fn agent_config() -> DQNConfig {
        DQNConfig {
            state_dim: 64, // TradingState::default() produces 4×16 = 64 features
            num_actions: 3,
            hidden_dims: vec![128, 64, 32],
            batch_size: 32,
            replay_buffer_capacity: 100_000,
            ..DQNConfig::default()
        }
    }

    #[tokio::test]
    async fn test_dqn_agent_creation() -> Result<(), Box<dyn std::error::Error>> {
        let config = agent_config();
        let agent = DQNAgent::new(config)?;

        assert_eq!(agent.get_epsilon(), 0.0); // Noisy networks: epsilon always 0
        assert_eq!(agent.get_config().state_dim, 64);

        Ok(())
    }

    #[test]
    fn test_trading_action_conversion() {
        assert_eq!(TradingAction::Buy.to_int(), 0);
        assert_eq!(TradingAction::Sell.to_int(), 1);
        assert_eq!(TradingAction::Hold.to_int(), 2);

        assert_eq!(TradingAction::from_int(0), Some(TradingAction::Buy));
        assert_eq!(TradingAction::from_int(1), Some(TradingAction::Sell));
        assert_eq!(TradingAction::from_int(2), Some(TradingAction::Hold));
        assert_eq!(TradingAction::from_int(3), None);
    }

    #[tokio::test]
    async fn test_action_selection() -> Result<(), Box<dyn std::error::Error>> {
        let config = agent_config();
        let mut agent = DQNAgent::new(config)?;
        let state = TradingState::default();

        let action = agent.select_action(&state)?;

        // Should be one of the three valid actions
        assert!(matches!(
            action,
            TradingAction::Buy | TradingAction::Sell | TradingAction::Hold
        ));

        Ok(())
    }

    #[tokio::test]
    async fn test_experience_storage() -> Result<(), Box<dyn std::error::Error>> {
        let config = agent_config();
        let mut agent = DQNAgent::new(config)?;

        let experience = Experience::new(
            vec![1.0; 64],                     // state (match default state_dim)
            TradingAction::Buy.to_int() as u8, // action
            100.0,                             // reward
            vec![1.1; 64],                     // next_state (match default state_dim)
            false,                             // done
        );

        agent.store_experience(experience)?;
        assert_eq!(agent.replay_buffer.size(), 1);

        Ok(())
    }

    #[tokio::test]
    async fn test_training_readiness() -> Result<(), Box<dyn std::error::Error>> {
        let config = agent_config();
        let mut agent = DQNAgent::new(config)?;

        // Should not be ready initially
        assert!(!agent.is_ready_for_training());

        // Add enough experiences
        for i in 0..400 {
            // More than batch_size * 10
            let experience = Experience::new(
                vec![i as f32; 64],
                TradingAction::Hold.to_int() as u8,
                i as f32,
                vec![i as f32 + 0.1; 64],
                i % 100 == 0, // Some terminal states
            );
            agent.store_experience(experience)?;
        }

        // Should be ready now
        assert!(agent.is_ready_for_training());

        Ok(())
    }

    #[tokio::test]
    async fn test_training_statistics() -> Result<(), Box<dyn std::error::Error>> {
        let config = agent_config();
        let mut agent = DQNAgent::new(config)?;

        // Update some statistics
        agent.update_reward_stats(100.0, true);
        agent.update_reward_stats(50.0, false);
        agent.update_reward_stats(75.0, true);

        let stats = agent.get_training_stats();
        assert_eq!(stats["total_episodes"], 3.0);
        assert!(stats["avg_reward"] > 0.0);
        assert!(stats["win_rate"] > 0.0 && stats["win_rate"] < 1.0);

        Ok(())
    }

    #[tokio::test]
    async fn test_network_summary() -> Result<(), Box<dyn std::error::Error>> {
        let config = agent_config();
        let agent = DQNAgent::new(config)?;

        let summary = agent.get_network_summary();
        assert!(summary.contains("DQN Network"));
        assert!(summary.contains("State Dimension: 64"));
        assert!(summary.contains("Action Space: 3"));

        Ok(())
    }

    #[test]
    fn test_trading_state_creation_and_validation() {
        use common::types::Price;
        use rust_decimal::Decimal;

        let state = TradingState::new(
            vec![
                Price::from_f64(1.0).unwrap(),
                Price::from_f64(2.0).unwrap(),
                Price::from_f64(3.0).unwrap(),
            ],
            vec![0.5, 0.6],
            vec![0.1, 0.2, 0.3, 0.4],
            vec![
                Decimal::try_from(100.0).unwrap(),
                Decimal::try_from(200.0).unwrap(),
            ],
        );

        assert!(state.is_valid());
        assert_eq!(state.dimension(), 11);

        let vector = state.to_vector();
        assert_eq!(vector.len(), 11);
        assert_eq!(vector[0], 1.0);
        assert_eq!(vector[3], 0.5);
        assert_eq!(vector[5], 0.1);
        assert_eq!(vector[9], 100.0);
    }

    #[test]
    fn test_trading_state_invalid_cases() {
        use rust_decimal::Decimal;

        let invalid_state = TradingState::new(
            vec![], // Empty price features should make it invalid
            vec![0.5],
            vec![0.1],
            vec![Decimal::try_from(100.0).unwrap()],
        );

        assert!(!invalid_state.is_valid());
    }

    #[test]
    fn test_trading_action_all() {
        let all_actions = TradingAction::all();
        assert_eq!(all_actions.len(), 3);
        assert_eq!(all_actions[0], TradingAction::Buy);
        assert_eq!(all_actions[1], TradingAction::Sell);
        assert_eq!(all_actions[2], TradingAction::Hold);
    }

    #[test]
    fn test_agent_metrics_default() {
        use rust_decimal::Decimal;

        let metrics = AgentMetrics::default();
        assert_eq!(metrics.total_episodes, 0);
        assert_eq!(metrics.total_steps, 0);
        assert_eq!(metrics.epsilon, 1.0);
        assert_eq!(metrics.avg_reward, Decimal::ZERO);
        assert_eq!(metrics.win_rate, 0.0);
        assert_eq!(metrics.current_loss, Decimal::ZERO);
    }

    #[tokio::test]
    async fn test_dqn_config_custom() -> Result<(), Box<dyn std::error::Error>> {
        let config = DQNConfig {
            state_dim: 32,
            num_actions: 3,
            hidden_dims: vec![64, 32],
            learning_rate: 0.01,
            gamma: 0.95,
            replay_buffer_capacity: 50_000,
            batch_size: 64,
            target_update_freq: 500,
            epsilon_start: 0.9,
            epsilon_end: 0.05,
            epsilon_decay: 0.99,
            minimum_profit_factor: 1.5,
            tau: 0.005,
            weight_decay: 1e-4,
            ..DQNConfig::default()
        };

        let agent = DQNAgent::new(config.clone())?;
        assert_eq!(agent.get_config().state_dim, 32);
        assert_eq!(agent.get_config().batch_size, 64);
        assert_eq!(agent.get_config().gamma, 0.95);

        Ok(())
    }

    #[tokio::test]
    async fn test_parameter_count_estimation() -> Result<(), Box<dyn std::error::Error>> {
        let config = DQNConfig {
            state_dim: 10,
            hidden_dims: vec![5, 3],
            num_actions: 2,
            ..DQNConfig::default()
        };

        let agent = DQNAgent::new(config)?;
        let param_count = agent.estimate_parameter_count();

        // Layer 1: 10*5 + 5 = 55
        // Layer 2: 5*3 + 3 = 18
        // Output: 3*2 + 2 = 8
        // Total = 81, doubled for target network = 162
        assert_eq!(param_count, 162);

        Ok(())
    }
}
