#![allow(unsafe_code)] // Required for CUDA cuBLAS sgemm and kernel launches.

//! Noisy Networks for Deep Reinforcement Learning
//!
//! Implementation of factorized Gaussian noise for exploration
//! as described in "Noisy Networks for Exploration" (Fortunato et al., 2018)
//!
//! This replaces epsilon-greedy exploration with learnable parameter noise.
//!
//! Architecture:
//! - Learnable μ (mean) and σ (std dev) parameters for weights and biases
//! - Factorized Gaussian noise: `ε_ij` = `f(ε_i)` × `f(ε_j)` where f(x) = sign(x) × √|x|
//! - Reduces parameter count by ~70% vs independent noise while maintaining exploration quality

use std::mem::ManuallyDrop;
use std::sync::Arc;

use cudarc::cublas::CudaBlas;
use cudarc::cublas::sys::cublasOperation_t;
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use ml_core::cuda_autograd::GpuTensor;
use ml_core::MLError;

// ── Pointer helpers (same pattern as gpu_tensor.rs) ──────────────────────

/// Extract raw CUDA device pointer from a `CudaSlice` (read-only).
fn raw_ptr(slice: &CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = ManuallyDrop::new(guard);
    ptr
}

/// Extract raw mutable CUDA device pointer from a `CudaSlice`.
fn raw_ptr_mut(slice: &mut CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr_mut(stream);
    let _no_drop = ManuallyDrop::new(guard);
    ptr
}

/// Compile the bias-add kernel for `NoisyLinear` forward.
fn compile_bias_add_kernel(stream: &Arc<CudaStream>) -> Result<cudarc::driver::CudaFunction, MLError> {
    let src = r#"
extern "C" __global__
void noisy_add_bias_kernel(float* __restrict__ y,
                           const float* __restrict__ bias,
                           int rows, int cols) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < rows * cols) {
        int col = idx % cols;
        y[idx] += bias[col];
    }
}
"#;
    let context = stream.context();
    let ptx: Ptx = ml_core::cuda_compile::compile_ptx_for_device(src, context)
        .map_err(|e| MLError::ModelError(format!("noisy bias kernel compilation: {e}")))?;
    let module = context.load_module(ptx).map_err(|e| {
        MLError::ModelError(format!("noisy bias module load: {e}"))
    })?;
    module.load_function("noisy_add_bias_kernel").map_err(|e| {
        MLError::ModelError(format!("noisy_add_bias_kernel load: {e}"))
    })
}

/// Noisy linear layer with factorized Gaussian noise (Rainbow DQN standard)
///
/// Key features:
/// - Learnable exploration: gradients tune noise parameters for optimal exploration
/// - Factorized noise: O(n+m) vs O(n×m) parameters (memory efficient)
/// - No epsilon decay needed: exploration naturally anneals via gradient descent
#[derive(Debug)]
pub struct NoisyLinear {
    // Learnable mean parameters (equivalent to standard Linear layer)
    weight_mu: CudaSlice<f32>,
    bias_mu: CudaSlice<f32>,

    // Learnable noise std dev parameters
    weight_sigma: CudaSlice<f32>,
    bias_sigma: CudaSlice<f32>,

    // Noise buffers (resampled each forward pass, not learned)
    weight_epsilon: CudaSlice<f32>,
    bias_epsilon: CudaSlice<f32>,

    // Dimensions
    in_features: usize,
    out_features: usize,

    // CUDA stream for GPU operations
    stream: Arc<CudaStream>,
}

impl NoisyLinear {
    /// Create new noisy linear layer
    ///
    /// # Arguments
    /// * `in_features` - Input dimension
    /// * `out_features` - Output dimension
    /// * `stream` - CUDA stream for GPU operations
    /// * `sigma_init` - Initial noise std dev before scaling (Rainbow DQN default: 0.5)
    ///
    /// # Initialization (Rainbow DQN standard):
    /// - `mu_w` ~ U(-1/sqrt(in), 1/sqrt(in))  (uniform distribution)
    /// - `sigma_w` = `sigma_init` / sqrt(in)   (factorized noise std dev)
    /// - Same for biases
    pub fn new(
        in_features: usize,
        out_features: usize,
        stream: Arc<CudaStream>,
        sigma_init: f64,
    ) -> Result<Self, MLError> {
        let mu_range = 1.0 / (in_features as f64).sqrt();
        let sigma_init_val = (sigma_init / (in_features as f64).sqrt()) as f32;

        // Initialize weight_mu ~ U(-1/sqrt(in), 1/sqrt(in))
        let weight_mu_host: Vec<f32> = (0..out_features * in_features)
            .map(|_| {
                let r: f32 = rand::random::<f32>() * 2.0 - 1.0;
                r * mu_range as f32
            })
            .collect();
        let (weight_mu, _) = GpuTensor::from_host(&weight_mu_host, vec![out_features, in_features], &stream)
            .map_err(|e| MLError::ModelError(format!("Failed to init weight_mu: {e}")))?
            .into_parts();

        // Initialize sigma_w = sigma_init / sqrt(in)
        let weight_sigma_host = vec![sigma_init_val; out_features * in_features];
        let (weight_sigma, _) = GpuTensor::from_host(&weight_sigma_host, vec![out_features, in_features], &stream)
            .map_err(|e| MLError::ModelError(format!("Failed to init weight_sigma: {e}")))?
            .into_parts();

        // Initialize bias mu and sigma
        let bias_mu_host: Vec<f32> = (0..out_features)
            .map(|_| {
                let r: f32 = rand::random::<f32>() * 2.0 - 1.0;
                r * mu_range as f32
            })
            .collect();
        let (bias_mu, _) = GpuTensor::from_host(&bias_mu_host, vec![out_features], &stream)
            .map_err(|e| MLError::ModelError(format!("Failed to init bias_mu: {e}")))?
            .into_parts();

        let bias_sigma_host = vec![sigma_init_val; out_features];
        let (bias_sigma, _) = GpuTensor::from_host(&bias_sigma_host, vec![out_features], &stream)
            .map_err(|e| MLError::ModelError(format!("Failed to init bias_sigma: {e}")))?
            .into_parts();

        // Zero-init epsilon buffers
        let weight_epsilon = stream.alloc_zeros::<f32>(out_features * in_features)
            .map_err(|e| MLError::ModelError(format!("Failed to init weight_epsilon: {e}")))?;
        let bias_epsilon = stream.alloc_zeros::<f32>(out_features)
            .map_err(|e| MLError::ModelError(format!("Failed to init bias_epsilon: {e}")))?;

        Ok(Self {
            weight_mu,
            bias_mu,
            weight_sigma,
            bias_sigma,
            weight_epsilon,
            bias_epsilon,
            in_features,
            out_features,
            stream,
        })
    }

    /// Resample noise (call before each forward pass during training)
    ///
    /// Uses factorized Gaussian noise: `epsilon_ij` = `f(epsilon_i)` x `f(epsilon_j)`
    /// where f(x) = sign(x) x sqrt(|x|) (reduces correlation)
    pub fn reset_noise(&mut self) -> Result<(), MLError> {
        self.reset_noise_with_sigma(1.0)
    }

    /// Resample noise with custom sigma scaling (for annealing)
    pub fn reset_noise_with_sigma(&mut self, sigma_scale: f64) -> Result<(), MLError> {
        // Factorized Gaussian noise: epsilon_ij = f(epsilon_i) * f(epsilon_j)
        // where f(x) = sign(x) * sqrt(|x|)
        let scale = sigma_scale as f32;

        // Generate raw Gaussian noise on host, apply f(x) = sign(x) * sqrt(|x|)
        let apply_f = |x: f32| -> f32 {
            let sign = if x >= 0.0 { 1.0_f32 } else { -1.0_f32 };
            sign * x.abs().sqrt()
        };

        // Input noise: p = in_features
        let eps_p: Vec<f32> = (0..self.in_features)
            .map(|_| {
                let r: f32 = rand::random::<f32>() * 2.0 - 1.0;
                // Box-Muller approximation: good enough for noise generation
                let n = r * std::f32::consts::SQRT_2; // rough normal approximation
                apply_f(n) * scale
            })
            .collect();

        // Output noise: q = out_features
        let eps_q: Vec<f32> = (0..self.out_features)
            .map(|_| {
                let r: f32 = rand::random::<f32>() * 2.0 - 1.0;
                let n = r * std::f32::consts::SQRT_2;
                apply_f(n) * scale
            })
            .collect();

        // Weight noise: outer product eps_q x eps_p -> [out_features, in_features]
        let mut weight_noise = Vec::with_capacity(self.out_features * self.in_features);
        for q_val in &eps_q {
            for p_val in &eps_p {
                weight_noise.push(q_val * p_val);
            }
        }

        // Bias noise: eps_q
        let bias_noise = eps_q;

        // Upload to GPU
        self.stream.memcpy_htod(&weight_noise, &mut self.weight_epsilon).map_err(|e| {
            MLError::ModelError(format!("reset_noise weight_epsilon HtoD: {e}"))
        })?;
        self.stream.memcpy_htod(&bias_noise, &mut self.bias_epsilon).map_err(|e| {
            MLError::ModelError(format!("reset_noise bias_epsilon HtoD: {e}"))
        })?;

        Ok(())
    }

    /// Forward pass with noisy weights (cold path -- tests and weight init only).
    ///
    /// **Hot-path forward is handled by `GpuDqnTrainer::forward_only_q()` which
    /// uses the fused CUDA kernel `dqn_forward_only_kernel` with BF16 tensor core
    /// matmul.**
    ///
    /// Computes entirely on GPU: `y = (mu_w + sigma_w * epsilon_w) x input + (mu_b + sigma_b * epsilon_b)`
    /// Zero CPU downloads -- all arithmetic via CUDA elementwise kernels + cuBLAS sgemm.
    #[cold]
    pub fn forward(&self, x: &GpuTensor) -> Result<GpuTensor, MLError> {
        let batch = if x.ndim() == 1 { 1 } else {
            x.shape().first().copied().unwrap_or(1)
        };
        let in_dim = self.in_features;
        let out_dim = self.out_features;
        let w_shape = vec![out_dim, in_dim];
        let b_shape = vec![out_dim];

        // Wrap CudaSlice refs as GpuTensor (cheap clone — cudarc 0.19 ref-counted)
        let mu_w = GpuTensor::new(self.weight_mu.clone(), w_shape.clone())?;
        let sigma_w = GpuTensor::new(self.weight_sigma.clone(), w_shape)?;
        let eps_w = GpuTensor::new(self.weight_epsilon.clone(), vec![out_dim, in_dim])?;

        let mu_b = GpuTensor::new(self.bias_mu.clone(), b_shape.clone())?;
        let sigma_b = GpuTensor::new(self.bias_sigma.clone(), b_shape)?;
        let eps_b = GpuTensor::new(self.bias_epsilon.clone(), vec![out_dim])?;

        // GPU elementwise: effective_weight = mu_w + sigma_w * epsilon_w
        let sigma_eps_w = sigma_w.mul(&eps_w, &self.stream)?;
        let effective_w = mu_w.add(&sigma_eps_w, &self.stream)?;

        // GPU elementwise: effective_bias = mu_b + sigma_b * epsilon_b
        let sigma_eps_b = sigma_b.mul(&eps_b, &self.stream)?;
        let effective_b = mu_b.add(&sigma_eps_b, &self.stream)?;

        // cuBLAS sgemm: y = x @ W_eff^T  (x:[batch, in], W_eff:[out, in] -> y:[batch, out])
        // Column-major: Y_col[out, B] = W_col^T[out, in] @ X_col[in, B]
        //   transA=T, transB=N, m=out, n=B, k=in
        let cublas = CudaBlas::new(self.stream.clone()).map_err(|e| {
            MLError::ModelError(format!("NoisyLinear forward cublas init: {e}"))
        })?;

        let mut y = GpuTensor::zeros(&[batch, out_dim], &self.stream)?;
        let w_ptr = raw_ptr(effective_w.data(), &self.stream);
        let x_ptr = raw_ptr(x.data(), &self.stream);
        let y_ptr = raw_ptr_mut(y.data_mut(), &self.stream);

        // SAFETY: w_ptr, x_ptr, y_ptr are valid device allocations on same context.
        // Dimensions are consistent: W[out, in] x X^T[in, batch] -> Y[out, batch].
        unsafe {
            cudarc::cublas::result::sgemm(
                *cublas.handle(),
                cublasOperation_t::CUBLAS_OP_T,    // transA: W stored [out, in], need W^T
                cublasOperation_t::CUBLAS_OP_N,    // transB: X stored [batch, in] = X_col[in, batch]
                out_dim as i32,                     // m
                batch as i32,                       // n
                in_dim as i32,                      // k
                &1.0_f32 as *const f32,            // alpha
                w_ptr as *const f32,               // A = W_eff
                in_dim as i32,                      // lda
                x_ptr as *const f32,               // B = X
                in_dim as i32,                      // ldb
                &0.0_f32 as *const f32,            // beta
                y_ptr as *mut f32,                 // C = Y
                out_dim as i32,                     // ldc
            ).map_err(|e| MLError::ModelError(format!("NoisyLinear cuBLAS sgemm: {e:?}")))?;
        }

        // Add bias: Y[b, j] += effective_bias[j] via CUDA kernel
        let bias_fn = compile_bias_add_kernel(&self.stream)?;
        let total = batch * out_dim;
        let threads = 256_u32;
        let blocks = (total as u32).div_ceil(threads);
        let rows_i32 = batch as i32;
        let cols_i32 = out_dim as i32;

        let launch_cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: 0,
        };

        // SAFETY: y and effective_b are valid device allocations, kernel indices are bounds-checked.
        unsafe {
            self.stream
                .launch_builder(&bias_fn)
                .arg(y.data())
                .arg(effective_b.data())
                .arg(&rows_i32)
                .arg(&cols_i32)
                .launch(launch_cfg)
                .map_err(|e| MLError::ModelError(format!("NoisyLinear bias kernel: {e}")))?;
        }

        Ok(y)
    }

    /// Get all learnable parameters as `CudaSlice` references (for optimizer)
    pub fn param_slices(&self) -> Vec<&CudaSlice<f32>> {
        vec![&self.weight_mu, &self.bias_mu, &self.weight_sigma, &self.bias_sigma]
    }

    /// Get only mu (mean) parameter slices -- `weight_mu`, `bias_mu`.
    pub const fn mu_slices(&self) -> [&CudaSlice<f32>; 2] {
        [&self.weight_mu, &self.bias_mu]
    }

    /// Input dimension of this layer.
    pub const fn in_features(&self) -> usize {
        self.in_features
    }

    /// Output dimension of this layer.
    pub const fn out_features(&self) -> usize {
        self.out_features
    }

    /// Get only sigma (noise std dev) parameter slices -- `weight_sigma`, `bias_sigma`.
    pub const fn sigma_slices(&self) -> [&CudaSlice<f32>; 2] {
        [&self.weight_sigma, &self.bias_sigma]
    }

    /// Disable noise for evaluation (use mean parameters only)
    pub fn disable_noise(&mut self) -> Result<(), MLError> {
        self.weight_epsilon = self.stream.alloc_zeros::<f32>(self.out_features * self.in_features)
            .map_err(|e| MLError::ModelError(format!("Failed to zero weight_epsilon: {e}")))?;
        self.bias_epsilon = self.stream.alloc_zeros::<f32>(self.out_features)
            .map_err(|e| MLError::ModelError(format!("Failed to zero bias_epsilon: {e}")))?;
        Ok(())
    }

    /// All data is already F32 on GPU -- no-op (Candle `DType` conversion removed).
    pub const fn ensure_f32(&mut self) -> Result<(), MLError> {
        // All CudaSlice<f32> are natively F32 -- nothing to convert.
        Ok(())
    }

    /// Copy all learnable parameters from `src` to `self` via GPU `DtoD` memcpy.
    ///
    /// Copies `weight_mu`, `bias_mu`, `weight_sigma`, `bias_sigma`, and epsilon buffers.
    /// Zero CPU involvement — all copies are async device-to-device.
    pub fn copy_params_from(&mut self, src: &NoisyLinear) -> Result<(), MLError> {
        Self::dtod_copy(&src.weight_mu, &mut self.weight_mu, "weight_mu", &self.stream)?;
        Self::dtod_copy(&src.bias_mu, &mut self.bias_mu, "bias_mu", &self.stream)?;
        Self::dtod_copy(&src.weight_sigma, &mut self.weight_sigma, "weight_sigma", &self.stream)?;
        Self::dtod_copy(&src.bias_sigma, &mut self.bias_sigma, "bias_sigma", &self.stream)?;
        Self::dtod_copy(&src.weight_epsilon, &mut self.weight_epsilon, "weight_epsilon", &self.stream)?;
        Self::dtod_copy(&src.bias_epsilon, &mut self.bias_epsilon, "bias_epsilon", &self.stream)?;
        Ok(())
    }

    /// Async device-to-device memcpy for a single `CudaSlice` buffer.
    fn dtod_copy(
        src: &CudaSlice<f32>,
        dst: &mut CudaSlice<f32>,
        label: &str,
        stream: &Arc<CudaStream>,
    ) -> Result<(), MLError> {
        let num_bytes = src.len() * std::mem::size_of::<f32>();
        let src_ptr = raw_ptr(src, stream);
        let dst_ptr = raw_ptr_mut(dst, stream);
        // SAFETY: src and dst are valid device allocations on the same context.
        // num_bytes = len * sizeof(f32) does not exceed either allocation.
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
            ).map_err(|e| MLError::ModelError(format!("NoisyLinear DtoD {label}: {e}")))?;
        }
        Ok(())
    }
}

/// Configuration for noisy networks
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct NoisyNetworkConfig {
    /// Initial noise std dev (Rainbow DQN: 0.5 / √`in_features`)
    pub sigma_init: f64,
    /// Whether to use noisy networks (replaces epsilon-greedy)
    pub enabled: bool,
}

impl Default for NoisyNetworkConfig {
    fn default() -> Self {
        Self {
            sigma_init: 0.5, // Rainbow DQN standard (scaled by 1/√in during init)
            enabled: false,  // Disabled by default (use epsilon-greedy)
        }
    }
}

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

    fn make_stream() -> Arc<CudaStream> {
        let ctx = cudarc::driver::CudaContext::new(0).ok();
        ctx.and_then(|c| c.new_stream().ok())
            .unwrap_or_else(|| panic!("CUDA stream required for NoisyLinear tests"))
    }

    #[test]
    fn test_noisy_linear_creation() -> Result<(), MLError> {
        let stream = make_stream();
        let _layer = NoisyLinear::new(64, 32, stream, 0.5)?;
        Ok(())
    }

    #[test]
    fn test_disable_noise() -> Result<(), MLError> {
        let stream = make_stream();
        let mut layer = NoisyLinear::new(64, 32, stream, 0.5)?;
        layer.disable_noise()?;
        // After disable, epsilon buffers should have the right length
        assert_eq!(layer.weight_epsilon.len(), 64 * 32);
        assert_eq!(layer.bias_epsilon.len(), 32);
        Ok(())
    }

    #[test]
    fn test_factorized_noise_dimensions() -> Result<(), MLError> {
        let stream = make_stream();
        let layer = NoisyLinear::new(128, 64, stream, 0.5)?;
        assert_eq!(layer.weight_epsilon.len(), 64 * 128);
        assert_eq!(layer.bias_epsilon.len(), 64);
        Ok(())
    }

    #[test]
    fn test_ensure_f32_is_noop() -> Result<(), MLError> {
        let stream = make_stream();
        let mut layer = NoisyLinear::new(32, 16, stream, 0.5)?;
        // Should be a no-op since all data is already f32
        layer.ensure_f32()?;
        Ok(())
    }
}
