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

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};

use cudarc::cublas::CudaBlas;
use cudarc::driver::CudaStream;

use ml_core::cuda_autograd::{
    ActivationKernels, GpuDropout, GpuLinear, GpuTensor, GpuVarStore,
};
use ml_core::device::MlDevice;
use ml_core::MLError;

/// Adaptive dropout scheduler that decreases dropout rate over training
/// (Wave 26 P1.6)
#[derive(Debug, Clone)]
pub struct DropoutScheduler {
    initial_rate: f64,
    final_rate: f64,
    decay_steps: usize,
    current_step: usize,
}

impl DropoutScheduler {
    /// Create a new dropout scheduler
    ///
    /// # Arguments
    /// * `initial_rate` - Starting dropout rate (e.g., 0.5 for high regularization)
    /// * `final_rate` - Ending dropout rate (e.g., 0.1 for fine-tuning)
    /// * `decay_steps` - Number of steps over which to decay
    pub const fn new(initial_rate: f64, final_rate: f64, decay_steps: usize) -> Self {
        Self {
            initial_rate,
            final_rate,
            decay_steps,
            current_step: 0,
        }
    }

    /// Get current dropout rate based on training progress
    pub fn get_rate(&self) -> f64 {
        if self.decay_steps == 0 {
            return self.final_rate;
        }

        let progress = (self.current_step as f64 / self.decay_steps as f64).min(1.0);
        self.initial_rate * (1.0 - progress) + self.final_rate * progress
    }

    /// Advance the scheduler by a number of steps
    pub const fn step(&mut self, steps: usize) {
        self.current_step = self.current_step.saturating_add(steps);
    }

    /// Get current step count
    pub const fn current_step(&self) -> usize {
        self.current_step
    }
}

/// Configuration for Q-Network
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
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,
    /// Target network update frequency
    pub target_update_freq: usize,
    /// Dropout probability (static, used when `dropout_schedule` is None)
    pub dropout_prob: f64,
    /// Adaptive dropout schedule: (`initial_rate`, `final_rate`, `decay_steps`)
    /// High dropout early prevents overfitting, low dropout late allows fine-tuning
    /// (Wave 26 P1.6)
    pub dropout_schedule: Option<(f64, f64, usize)>,
    /// Whether to use `GPU` acceleration
    pub use_gpu: bool,
    /// Whether to use spectral normalization for Q-value stability
    pub use_spectral_norm: bool,
    /// Number of power iterations for spectral norm estimation
    pub spectral_norm_iterations: usize,
    /// Whether to use residual connections (Wave 26 P0.4)
    pub use_residual: 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,
            target_update_freq: 1000,
            dropout_prob: 0.2,
            dropout_schedule: None, // No adaptive dropout by default
            use_gpu: true,
            use_spectral_norm: false, // Disabled by default, enable for unstable training
            spectral_norm_iterations: 1, // 1 iteration is typically sufficient
            use_residual: false, // Disabled by default, opt-in for deeper networks (Wave 26 P0.4)

        }
    }
}

/// Deep Q-Network implementation
#[allow(missing_debug_implementations)]
pub struct QNetwork {
    /// Network configuration
    config: QNetworkConfig,
    /// Main network variables
    vars: GpuVarStore,
    /// Target network variables
    target_vars: GpuVarStore,
    /// Main network linear layers (names referencing [`vars`])
    layers: Vec<GpuLinear>,
    /// Target network linear layers (names referencing `target_vars`)
    _target_layers: Vec<GpuLinear>,
    /// CUDA stream for compute
    stream: Arc<CudaStream>,
    /// cuBLAS handle for matmul
    cublas: CudaBlas,
    /// Activation kernels
    activations: ActivationKernels,
    /// Dropout layer
    dropout: std::sync::Mutex<GpuDropout>,
    /// Training step counter
    step_count: AtomicU64,
    /// Adaptive dropout scheduler (Wave 26 P1.6)
    dropout_scheduler: std::sync::Mutex<Option<DropoutScheduler>>,
    /// Whether the network is in training mode (enables dropout)
    training: AtomicBool,
}

impl QNetwork {
    /// Create a new Q-Network
    pub fn new(config: QNetworkConfig) -> Result<Self, MLError> {
        if !config.use_gpu {
            return Err(MLError::DeviceError("CUDA required -- set use_gpu=true".into()));
        }

        let device = MlDevice::cuda(0)?;
        let stream = device.cuda_stream()?.clone();

        let cublas = CudaBlas::new(stream.clone()).map_err(|e| {
            MLError::ModelError(format!("Failed to create cuBLAS handle: {e}"))
        })?;
        let activations = ActivationKernels::new(&stream)?;

        // Build main network layers
        let mut vars = GpuVarStore::new(stream.clone());
        let layers = Self::build_layers(&mut vars, &config, "main")?;

        // Build target network layers (same architecture)
        let mut target_vars = GpuVarStore::new(stream.clone());
        let target_layers = Self::build_layers(&mut target_vars, &config, "target")?;

        // Initialize dropout scheduler if configured (Wave 26 P1.6)
        let dropout_scheduler = config.dropout_schedule.map(
            |(initial, final_rate, steps)| DropoutScheduler::new(initial, final_rate, steps),
        );

        let dropout = std::sync::Mutex::new(GpuDropout::new(config.dropout_prob as f32));

        Ok(Self {
            config,
            vars,
            target_vars,
            layers,
            _target_layers: target_layers,
            stream,
            cublas,
            activations,
            dropout,
            step_count: AtomicU64::new(0),
            dropout_scheduler: std::sync::Mutex::new(dropout_scheduler),
            training: AtomicBool::new(false),
        })
    }

    /// Build linear layers for a network, registering parameters in the given var store.
    fn build_layers(
        store: &mut GpuVarStore,
        config: &QNetworkConfig,
        prefix: &str,
    ) -> Result<Vec<GpuLinear>, MLError> {
        let mut layers = Vec::new();
        let mut input_dim = config.state_dim;

        // Hidden layers with Xavier initialization
        for (i, &hidden_dim) in config.hidden_dims.iter().enumerate() {
            let layer = store.linear(&format!("{prefix}.layer_{i}"), input_dim, hidden_dim)?;
            layers.push(layer);
            input_dim = hidden_dim;
        }

        // Output layer
        let output_layer = store.linear(&format!("{prefix}.output"), input_dim, config.num_actions)?;
        layers.push(output_layer);

        Ok(layers)
    }

    /// 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()
            )));
        }

        // Upload state to GPU as [1, state_dim]
        let mut x = GpuTensor::from_host(
            state,
            vec![1, self.config.state_dim],
            &self.stream,
        )?;

        let is_training = self.is_training();

        // Forward through hidden layers with LeakyReLU + dropout
        let layer_count = self.layers.len();
        for (i, layer) in self.layers.iter().enumerate() {
            let (output, _acts) = layer.forward(&x, &self.vars, &self.cublas, &self.stream)?;
            x = output;

            // Apply LeakyReLU + dropout for all layers except the last
            if i < layer_count - 1 {
                let (activated, _saved) = self.activations.leaky_relu_fwd(&x, 0.01, &self.stream)?;
                x = activated;
                if is_training {
                    if let Ok(mut dropout) = self.dropout.lock() {
                        let (dropped, _mask) = dropout.forward(&x, &self.stream)?;
                        x = dropped;
                    }
                }
            }
        }

        // Download output to CPU
        let output_vec = x.to_host(&self.stream)?;

        // Update step count
        let _step = self.step_count.fetch_add(1, Ordering::Relaxed);

        // Update dropout scheduler (Wave 26 P1.6)
        if let Ok(mut scheduler_opt) = self.dropout_scheduler.lock() {
            if let Some(scheduler) = scheduler_opt.as_mut() {
                scheduler.step(1);
            }
        }

        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);
        }

        // Upload to GPU as [batch_size, state_dim]
        let mut x = GpuTensor::from_host(
            &flat_states,
            vec![batch_size, state_dim],
            &self.stream,
        )?;

        let is_training = self.is_training();
        let layer_count = self.layers.len();

        for (i, layer) in self.layers.iter().enumerate() {
            let (output, _acts) = layer.forward(&x, &self.vars, &self.cublas, &self.stream)?;
            x = output;

            if i < layer_count - 1 {
                let (activated, _saved) = self.activations.leaky_relu_fwd(&x, 0.01, &self.stream)?;
                x = activated;
                if is_training {
                    if let Ok(mut dropout) = self.dropout.lock() {
                        let (dropped, _mask) = dropout.forward(&x, &self.stream)?;
                        x = dropped;
                    }
                }
            }
        }

        // Download and reshape into Vec<Vec<f32>>
        let flat_output = x.to_host(&self.stream)?;
        let num_actions = self.config.num_actions;
        let mut result = Vec::with_capacity(batch_size);
        for b in 0..batch_size {
            let start = b * num_actions;
            let end = start + num_actions;
            result.push(flat_output.get(start..end)
                .ok_or_else(|| MLError::ModelError("Output slice out of bounds".into()))?
                .to_vec()); // cpu-side: shape metadata
        }

        Ok(result)
    }

    /// Select action using greedy policy (exploration handled by noisy networks)
    pub fn select_action(&self, state: &[f32]) -> Result<usize, MLError> {
        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 device information
    pub fn device_info(&self) -> String {
        "CUDA".to_owned()
    }

    /// Get reference to the CUDA stream
    pub const fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

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

    /// Get mutable reference to the variables (for checkpoint loading)
    pub const fn vars_mut(&mut self) -> &mut GpuVarStore {
        &mut self.vars
    }

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

    /// Get mutable reference to the target variables (for checkpoint loading)
    pub const fn target_vars_mut(&mut self) -> &mut GpuVarStore {
        &mut self.target_vars
    }

    /// Get current dropout rate (adaptive if scheduler is configured, otherwise static)
    /// (Wave 26 P1.6)
    pub fn get_dropout_rate(&self) -> f64 {
        if let Ok(scheduler_opt) = self.dropout_scheduler.lock() {
            if let Some(scheduler) = scheduler_opt.as_ref() {
                return scheduler.get_rate();
            }
        }
        // Fallback to static dropout probability
        self.config.dropout_prob
    }

    /// Set training mode (enables/disables dropout)
    pub fn set_training(&self, training: bool) {
        self.training.store(training, Ordering::Relaxed);
        if let Ok(mut dropout) = self.dropout.lock() {
            dropout.training = training;
        }
    }

    /// Check if network is in training mode
    pub fn is_training(&self) -> bool {
        self.training.load(Ordering::Relaxed)
    }
}

// Use QNetwork directly - no compatibility wrapper needed

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


    #[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: 9,
            ..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(), 9); // 9 exposure levels (default)
        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(())
    }

}
