//! GPU dropout layer for regularization.
//!
//! At inference (default): identity pass-through, returns input unchanged.
//! At training: generates a random mask and multiplies element-wise, scaling
//! surviving activations by `1/(1-p)` to maintain expected values.
//!
//! Uses a CUDA kernel compiled via NVRTC for the training path.

use std::sync::Arc;

use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;

use crate::MLError;
use super::gpu_tensor::GpuTensor;
use super::linear::clone_gpu_tensor;

/// GPU dropout layer.
///
/// By default operates in inference mode (pass-through). Set `training = true`
/// to enable stochastic dropout during forward passes.
pub struct GpuDropout {
    /// Dropout probability (fraction of elements to zero out). Range: [0, 1).
    pub p: f32,
    /// Whether the layer is in training mode.
    pub training: bool,
    /// Compiled dropout kernel (lazily created on first training forward).
    kernel: Option<CudaFunction>,
    /// RNG seed for deterministic dropout masks.
    seed: u64,
}

impl std::fmt::Debug for GpuDropout {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "GpuDropout {{ p: {}, training: {} }}", self.p, self.training)
    }
}

impl GpuDropout {
    /// Create a new dropout layer with the given probability.
    ///
    /// Starts in inference mode (pass-through).
    pub fn new(p: f32) -> Self {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};
        use std::time::SystemTime;

        let mut hasher = DefaultHasher::new();
        SystemTime::now().hash(&mut hasher);
        std::thread::current().id().hash(&mut hasher);
        let seed = hasher.finish();

        Self {
            p,
            training: false,
            kernel: None,
            seed,
        }
    }

    /// Forward pass.
    ///
    /// - **Inference** (`training = false`): returns input unchanged.
    /// - **Training** (`training = true`): applies stochastic dropout mask.
    ///
    /// Returns `(output, mask)` where mask is `None` in inference mode.
    pub fn forward(
        &mut self,
        x: &GpuTensor,
        stream: &Arc<CudaStream>,
    ) -> Result<(GpuTensor, Option<GpuTensor>), MLError> {
        if !self.training || self.p <= 0.0 {
            // Inference: identity pass-through
            let out = clone_gpu_tensor(x, stream)?;
            return Ok((out, None));
        }

        if self.p >= 1.0 {
            // Drop everything
            let out = GpuTensor::zeros(&x.shape, stream)?;
            let mask = GpuTensor::zeros(&x.shape, stream)?;
            return Ok((out, Some(mask)));
        }

        // Ensure kernel is compiled
        if self.kernel.is_none() {
            self.compile_kernel(stream)?;
        }

        let n = x.numel();
        let y = GpuTensor::zeros(&x.shape, stream)?;
        let mask = GpuTensor::zeros(&x.shape, stream)?;
        let n_i32 = n as i32;
        let scale = 1.0_f32 / (1.0_f32 - self.p);

        // Advance seed for this call
        self.seed = self.seed.wrapping_mul(6364136223846793005).wrapping_add(1);
        let seed_lo = self.seed as u32;
        let seed_hi = (self.seed >> 32) as u32;

        let threads = 256_u32;
        let blocks = ((n as u32) + threads - 1) / threads;
        let cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: 0,
        };

        let kernel = self.kernel.as_ref().ok_or_else(|| {
            MLError::ModelError("dropout kernel not compiled".to_owned())
        })?;

        unsafe {
            stream
                .launch_builder(kernel)
                .arg(&x.data)
                .arg(&y.data)
                .arg(&mask.data)
                .arg(&self.p)
                .arg(&scale)
                .arg(&seed_lo)
                .arg(&seed_hi)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("dropout_forward: {e}")))?;
        }

        Ok((y, Some(mask)))
    }

    /// Backward pass: `dx = dy * mask`.
    ///
    /// In inference mode, `mask` should be `None` and `dy` passes through unchanged.
    pub fn backward(
        &self,
        dy: &GpuTensor,
        mask: Option<&GpuTensor>,
        stream: &Arc<CudaStream>,
    ) -> Result<GpuTensor, MLError> {
        match mask {
            None => clone_gpu_tensor(dy, stream),
            Some(m) => {
                // dx = dy * mask (mask already includes the 1/(1-p) scaling)
                // We need an element-wise multiply kernel. Reuse a simple approach:
                // allocate output, launch multiply kernel.
                let n = dy.numel();
                let dx = GpuTensor::zeros(&dy.shape, stream)?;

                // Compile a simple multiply kernel
                let src = r#"
extern "C" __global__
void elemwise_mul(const float* __restrict__ a,
                  const float* __restrict__ b,
                  float* __restrict__ c,
                  int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        c[i] = a[i] * b[i];
    }
}
"#;
                let context = stream.context();
                let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(src, &context)
                    .map_err(|e| MLError::ModelError(format!("dropout backward kernel: {e}")))?;
                let module = context.load_module(ptx)
                    .map_err(|e| MLError::ModelError(format!("dropout backward module: {e}")))?;
                let mul_fn = module.load_function("elemwise_mul")
                    .map_err(|e| MLError::ModelError(format!("dropout backward load: {e}")))?;

                let n_i32 = n as i32;
                let threads = 256_u32;
                let blocks = ((n as u32) + threads - 1) / threads;
                let cfg = LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (threads, 1, 1),
                    shared_mem_bytes: 0,
                };

                unsafe {
                    stream
                        .launch_builder(&mul_fn)
                        .arg(&dy.data)
                        .arg(&m.data)
                        .arg(&dx.data)
                        .arg(&n_i32)
                        .launch(cfg)
                        .map_err(|e| MLError::ModelError(format!("dropout backward: {e}")))?;
                }

                Ok(dx)
            }
        }
    }

    /// Compile the dropout forward kernel.
    fn compile_kernel(&mut self, stream: &Arc<CudaStream>) -> Result<(), MLError> {
        let context = stream.context();
        let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(DROPOUT_CUDA_SRC, &context)
            .map_err(|e| MLError::ModelError(format!("dropout kernel compilation: {e}")))?;
        let module = context.load_module(ptx)
            .map_err(|e| MLError::ModelError(format!("dropout module load: {e}")))?;
        let kernel = module.load_function("dropout_forward")
            .map_err(|e| MLError::ModelError(format!("dropout_forward load: {e}")))?;
        self.kernel = Some(kernel);
        Ok(())
    }
}

// ── CUDA source for dropout kernel ──────────────────────────────────────

const DROPOUT_CUDA_SRC: &str = r#"
// Dropout forward kernel with built-in xoshiro128+ PRNG.
// Each thread generates its own random value from (seed + thread_idx).
// mask[i] = (rand > p) ? scale : 0
// y[i] = x[i] * mask[i]

extern "C" __global__
void dropout_forward(const float* __restrict__ x,
                     float* __restrict__ y,
                     float* __restrict__ mask,
                     float p,
                     float scale,
                     unsigned int seed_lo,
                     unsigned int seed_hi,
                     int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        // Simple hash-based RNG: mix seed with thread index
        unsigned int h = seed_lo ^ (unsigned int)i;
        h *= 2654435761u;  // Knuth multiplicative hash
        h ^= seed_hi;
        h *= 2246822519u;
        h ^= (h >> 13);
        h *= 3266489917u;
        h ^= (h >> 16);

        // Convert to float in [0, 1)
        float r = (float)(h & 0x00FFFFFFu) / (float)0x01000000u;

        float m = (r >= p) ? scale : 0.0f;
        mask[i] = m;
        y[i] = x[i] * m;
    }
}
"#;
