foxhunt/ml/src/dqn/network.rs

//! Q-Network implementation with target network and GPU acceleration

use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};

use candle_core::{DType, Device, Result as CandleResult, Tensor};
use candle_nn::Module;
use candle_nn::{ops::leaky_relu, Dropout, Linear, VarBuilder, VarMap};
use rand::prelude::*; // Replace common::rng with standard rand

use crate::dqn::xavier_init::linear_xavier; // Xavier initialization
use crate::MLError;

/// Configuration for Q-Network
#[derive(Debug, Clone)]
pub struct QNetworkConfig {
    /// Input state dimensions
    pub state_dim: usize,
    /// Number of possible actions
    pub num_actions: usize,
    /// Hidden layer sizes
    pub hidden_dims: Vec<usize>,
    /// Learning rate
    pub learning_rate: f64,
    /// Exploration epsilon start
    pub epsilon_start: f64,
    /// Exploration epsilon end
    pub epsilon_end: f64,
    /// Epsilon decay rate
    pub epsilon_decay: f64,
    /// Target network update frequency
    pub target_update_freq: usize,
    /// Dropout probability
    pub dropout_prob: f64,
    /// Whether to use `GPU` acceleration
    pub use_gpu: bool,
}

impl Default for QNetworkConfig {
    fn default() -> Self {
        Self {
            state_dim: 64,
            num_actions: 3,
            hidden_dims: vec![128, 64, 32],
            learning_rate: 0.001,
            epsilon_start: 1.0,
            epsilon_end: 0.01,
            epsilon_decay: 0.995,
            target_update_freq: 1000,
            dropout_prob: 0.2,
            use_gpu: false,
        }
    }
}

/// Deep Q-Network implementation
#[allow(missing_debug_implementations)]
pub struct QNetwork {
    /// Network configuration
    config: QNetworkConfig,
    /// Main network variables
    vars: VarMap,
    /// Target network variables
    target_vars: VarMap,
    /// Compute device
    device: Device,
    /// Current epsilon for exploration
    epsilon: AtomicU32, // Store as fixed-point u32
    /// Training step counter
    step_count: AtomicU64,
}

/// Network layer structure
#[derive(Debug)]
struct NetworkLayers {
    layers: Vec<Linear>,
    dropout: Dropout,
}

impl NetworkLayers {
    fn new(
        var_builder: &VarBuilder<'_>,
        config: &QNetworkConfig,
        _device: &Device,
    ) -> CandleResult<Self> {
        let mut layers = Vec::new();
        let mut input_dim = config.state_dim;

        // Create hidden layers with Xavier initialization
        for (i, &hidden_dim) in config.hidden_dims.iter().enumerate() {
            // Xavier uniform initialization for better gradient flow
            let layer = linear_xavier(
                input_dim,
                hidden_dim,
                var_builder.pp(format!("layer_{}", i)),
            )?;
            layers.push(layer);
            input_dim = hidden_dim;
        }

        // Output layer - also use Xavier initialization
        let output_layer = linear_xavier(input_dim, config.num_actions, var_builder.pp("output"))?;
        layers.push(output_layer);

        let dropout = Dropout::new(config.dropout_prob as f32);

        Ok(Self { layers, dropout })
    }
}

impl Module for NetworkLayers {
    fn forward(&self, xs: &Tensor) -> CandleResult<Tensor> {
        let mut x = xs.clone();

        // Forward through hidden layers with LeakyReLU activation and dropout
        // LeakyReLU prevents dead neurons (0.01 gradient for negative inputs vs 0 for ReLU)
        for (i, layer) in self.layers.iter().enumerate() {
            x = layer.forward(&x)?;

            // Apply LeakyReLU activation for all layers except the last
            if i < self.layers.len() - 1 {
                x = leaky_relu(&x, 0.01)?; // Bug #11 fix: LeakyReLU prevents gradient collapse
                x = self.dropout.forward(&x, false)?; // No dropout during inference
            }
        }

        Ok(x)
    }
}

impl QNetwork {
    /// Create a new Q-Network
    pub fn new(config: QNetworkConfig) -> Result<Self, MLError> {
        let device = if config.use_gpu && Device::cuda_if_available(0).is_ok() {
            Device::new_cuda(0)
                .map_err(|e| MLError::ModelError(format!("Failed to initialize CUDA: {}", e)))?
        } else {
            Device::Cpu
        };

        let vars = VarMap::new();
        let target_vars = VarMap::new();

        // Initialize network weights
        let var_builder = VarBuilder::from_varmap(&vars, DType::F32, &device);
        let _layers = NetworkLayers::new(&var_builder, &config, &device)
            .map_err(|e| MLError::ModelError(format!("Failed to create network layers: {}", e)))?;

        // Initialize target network with same architecture
        let target_var_builder = VarBuilder::from_varmap(&target_vars, DType::F32, &device);
        let _target_layers =
            NetworkLayers::new(&target_var_builder, &config, &device).map_err(|e| {
                MLError::ModelError(format!("Failed to create target network layers: {}", e))
            })?;

        let epsilon = (config.epsilon_start * 1_000_000.0) as u32; // Fixed-point representation

        Ok(Self {
            config,
            vars,
            target_vars,
            device,
            epsilon: AtomicU32::new(epsilon),
            step_count: AtomicU64::new(0),
        })
    }

    /// Forward pass through the network
    pub fn forward(&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 var_builder = VarBuilder::from_varmap(&self.vars, DType::F32, &self.device);
        let layers = NetworkLayers::new(&var_builder, &self.config, &self.device)
            .map_err(|e| MLError::ModelError(format!("Failed to create layers: {}", e)))?;

        let input = Tensor::from_vec(state.to_vec(), state.len(), &self.device)
            .map_err(|e| MLError::ModelError(format!("Failed to create input tensor: {}", e)))?
            .unsqueeze(0) // Add batch dimension
            .map_err(|e| MLError::ModelError(format!("Failed to add batch dimension: {}", e)))?;

        let output = layers
            .forward(&input)
            .map_err(|e| MLError::ModelError(format!("Forward pass failed: {}", e)))?;

        let output_vec = output
            .squeeze(0)
            .map_err(|e| MLError::ModelError(format!("Failed to squeeze output: {}", e)))?
            .to_vec1::<f32>()
            .map_err(|e| {
                MLError::ModelError(format!("Failed to convert output to vector: {}", e))
            })?;

        // Update step count and decay epsilon
        let step = self.step_count.fetch_add(1, Ordering::Relaxed);
        self.decay_epsilon(step);

        Ok(output_vec)
    }

    /// Forward pass for batch of states
    pub fn forward_batch(&self, states: &[Vec<f32>]) -> Result<Vec<Vec<f32>>, MLError> {
        if states.is_empty() {
            return Ok(Vec::new());
        }

        let batch_size = states.len();
        let state_dim = self.config.state_dim;

        // Flatten states into single vector
        let mut flat_states = Vec::with_capacity(batch_size * state_dim);
        for state in states {
            if state.len() != state_dim {
                return Err(MLError::InvalidInput(format!(
                    "State dimension mismatch: expected {}, got {}",
                    state_dim,
                    state.len()
                )));
            }
            flat_states.extend_from_slice(state);
        }

        let var_builder = VarBuilder::from_varmap(&self.vars, DType::F32, &self.device);
        let layers = NetworkLayers::new(&var_builder, &self.config, &self.device)
            .map_err(|e| MLError::ModelError(format!("Failed to create layers: {}", e)))?;

        let input = Tensor::from_vec(flat_states, (batch_size, state_dim), &self.device)
            .map_err(|e| MLError::ModelError(format!("Failed to create input tensor: {}", e)))?;

        let output = layers
            .forward(&input)
            .map_err(|e| MLError::ModelError(format!("Forward pass failed: {}", e)))?;

        let output_vec = output.to_vec2::<f32>().map_err(|e| {
            MLError::ModelError(format!("Failed to convert output to vector: {}", e))
        })?;

        Ok(output_vec)
    }

    /// Select action using epsilon-greedy policy
    pub fn select_action(&self, state: &[f32]) -> Result<usize, MLError> {
        let epsilon = self.get_epsilon();

        if thread_rng().gen::<f64>() < epsilon {
            // Random action
            Ok(thread_rng().gen_range(0..self.config.num_actions))
        } else {
            // Greedy action selection
            let q_values = self.forward(state)?;
            let best_action = q_values
                .iter()
                .enumerate()
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                .map(|(idx, _)| idx)
                .unwrap_or(0);

            Ok(best_action)
        }
    }

    /// Get current epsilon value
    pub fn get_epsilon(&self) -> f64 {
        let epsilon_fixed = self.epsilon.load(Ordering::Relaxed);
        epsilon_fixed as f64 / 1_000_000.0
    }

    /// Set epsilon value
    pub fn set_epsilon(&self, epsilon: f64) {
        let epsilon_fixed = (epsilon.clamp(0.0, 1.0) * 1_000_000.0) as u32;
        self.epsilon.store(epsilon_fixed, Ordering::Relaxed);
    }

    /// Decay epsilon based on step count
    fn decay_epsilon(&self, step: u64) {
        if step > 0 && step % 100 == 0 {
            let current_epsilon = self.get_epsilon();
            let new_epsilon =
                (current_epsilon * self.config.epsilon_decay).max(self.config.epsilon_end);
            self.set_epsilon(new_epsilon);
        }
    }

    /// Get device information
    pub fn device_info(&self) -> String {
        match &self.device {
            Device::Cpu => "CPU".to_string(),
            Device::Cuda(_) => format!("CUDA"),
            Device::Metal(_) => "Metal".to_string(),
        }
    }

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

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

    /// Get reference to the target variables
    pub fn target_vars(&self) -> &VarMap {
        &self.target_vars
    }
}

// Use QNetwork directly - no compatibility wrapper needed

#[cfg(test)]
mod tests {
    use super::*;
    // use crate::safe_operations; // DISABLED - module not found

    #[test]
    fn test_qnetwork_creation() -> anyhow::Result<()> {
        let config = QNetworkConfig::default();
        let network = QNetwork::new(config)
            .map_err(|e| anyhow::anyhow!("Failed to create QNetwork: {:?}", e))?;

        let info = network.device_info();
        assert!(!info.is_empty());
        Ok(())
    }

    #[test]
    fn test_forward_pass() -> anyhow::Result<()> {
        let config = QNetworkConfig {
            state_dim: 4,
            num_actions: 5,
            ..QNetworkConfig::default()
        };
        let network = QNetwork::new(config)
            .map_err(|e| anyhow::anyhow!("Failed to create QNetwork: {:?}", e))?;

        let state = vec![1.0, 2.0, 3.0, 4.0];
        let q_values = network
            .forward(&state)
            .map_err(|e| anyhow::anyhow!("Forward pass failed: {:?}", e))?;

        assert_eq!(q_values.len(), 5);
        Ok(())
    }

    #[test]
    fn test_action_selection() -> anyhow::Result<()> {
        // Test action selection validation
        let num_actions = 5;
        let action = 2; // Sample action
        assert!(action < num_actions);
        Ok(())
    }

    #[test]
    fn test_batch_processing() -> anyhow::Result<()> {
        // Test batch processing concepts
        let batch_size = 2;
        let state_dim = 3;
        assert!(batch_size > 0);
        assert!(state_dim > 0);
        Ok(())
    }

    #[test]
    fn test_epsilon_decay() -> anyhow::Result<()> {
        // Test epsilon decay concepts
        let initial_epsilon = 1.0;
        let decay_rate = 0.995;
        let later_epsilon = initial_epsilon * decay_rate;

        assert!(initial_epsilon > 0.8);
        assert!(later_epsilon <= initial_epsilon);
        Ok(())
    }
}