//! Element-wise activation functions as CUDA kernels.
//!
//! Each activation provides:
//! - `forward(x)` -> `(output, saved)` where saved is whatever backward needs
//! - `backward(dy, saved)` -> `dx`
//!
//! All operations are in-place where possible to minimize memory allocation.

use std::sync::Arc;

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

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

/// Compiled activation kernels -- cached per CUDA context.
#[allow(missing_debug_implementations)]
pub struct ActivationKernels {
    relu_forward: CudaFunction,
    relu_backward: CudaFunction,
    leaky_relu_forward: CudaFunction,
    leaky_relu_backward: CudaFunction,
    gelu_forward: CudaFunction,
    gelu_backward: CudaFunction,
    sigmoid_forward: CudaFunction,
    sigmoid_backward: CudaFunction,
    tanh_forward: CudaFunction,
    tanh_backward: CudaFunction,
}

/// Launch config helper for element-wise kernels.
fn elem_cfg(n: usize) -> LaunchConfig {
    let threads = 256_u32;
    let blocks = ((n as u32) + threads - 1) / threads;
    LaunchConfig {
        grid_dim: (blocks, 1, 1),
        block_dim: (threads, 1, 1),
        shared_mem_bytes: 0,
    }
}

impl ActivationKernels {
    /// Compile all activation kernels from a single CUDA source.
    pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();
        let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(ACTIVATION_CUDA_SRC, &context)
            .map_err(|e| MLError::ModelError(format!("activation kernel compilation: {e}")))?;
        let module = context.load_module(ptx).map_err(|e| {
            MLError::ModelError(format!("activation module load: {e}"))
        })?;

        let load = |name: &str| -> Result<CudaFunction, MLError> {
            module.load_function(name).map_err(|e| {
                MLError::ModelError(format!("activation kernel '{name}' load: {e}"))
            })
        };

        Ok(Self {
            relu_forward: load("relu_forward")?,
            relu_backward: load("relu_backward")?,
            leaky_relu_forward: load("leaky_relu_forward")?,
            leaky_relu_backward: load("leaky_relu_backward")?,
            gelu_forward: load("gelu_forward")?,
            gelu_backward: load("gelu_backward")?,
            sigmoid_forward: load("sigmoid_forward")?,
            sigmoid_backward: load("sigmoid_backward")?,
            tanh_forward: load("tanh_forward")?,
            tanh_backward: load("tanh_backward")?,
        })
    }

    // ── ReLU ──────────────────────────────────────────────────────────

    /// ReLU forward: `y = max(0, x)`. Returns (output, mask) where mask = x > 0.
    pub fn relu_fwd(
        &self,
        x: &GpuTensor,
        stream: &Arc<CudaStream>,
    ) -> Result<(GpuTensor, GpuTensor), MLError> {
        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 cfg = elem_cfg(n);
        unsafe {
            stream
                .launch_builder(&self.relu_forward)
                .arg(&x.data)
                .arg(&y.data)
                .arg(&mask.data)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("relu_forward: {e}")))?;
        }
        Ok((y, mask))
    }

    /// ReLU backward: `dx = dy * mask`.
    pub fn relu_bwd(
        &self,
        dy: &GpuTensor,
        mask: &GpuTensor,
        stream: &Arc<CudaStream>,
    ) -> Result<GpuTensor, MLError> {
        let n = dy.numel();
        let dx = GpuTensor::zeros(&dy.shape, stream)?;
        let n_i32 = n as i32;
        let cfg = elem_cfg(n);
        unsafe {
            stream
                .launch_builder(&self.relu_backward)
                .arg(&dy.data)
                .arg(&mask.data)
                .arg(&dx.data)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("relu_backward: {e}")))?;
        }
        Ok(dx)
    }

    // ── Leaky ReLU ───────────────────────────────────────────────────

    /// Leaky ReLU forward with negative slope `alpha` (default 0.01).
    pub fn leaky_relu_fwd(
        &self,
        x: &GpuTensor,
        alpha: f32,
        stream: &Arc<CudaStream>,
    ) -> Result<(GpuTensor, GpuTensor), MLError> {
        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 cfg = elem_cfg(n);
        unsafe {
            stream
                .launch_builder(&self.leaky_relu_forward)
                .arg(&x.data)
                .arg(&y.data)
                .arg(&mask.data)
                .arg(&alpha)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("leaky_relu_forward: {e}")))?;
        }
        Ok((y, mask))
    }

    /// Leaky ReLU backward.
    pub fn leaky_relu_bwd(
        &self,
        dy: &GpuTensor,
        mask: &GpuTensor,
        alpha: f32,
        stream: &Arc<CudaStream>,
    ) -> Result<GpuTensor, MLError> {
        let n = dy.numel();
        let dx = GpuTensor::zeros(&dy.shape, stream)?;
        let n_i32 = n as i32;
        let cfg = elem_cfg(n);
        unsafe {
            stream
                .launch_builder(&self.leaky_relu_backward)
                .arg(&dy.data)
                .arg(&mask.data)
                .arg(&dx.data)
                .arg(&alpha)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("leaky_relu_backward: {e}")))?;
        }
        Ok(dx)
    }

    // ── GELU ─────────────────────────────────────────────────────────

    /// GELU forward (approximate: `0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))`).
    ///
    /// Saves the input `x` for the backward pass.
    pub fn gelu_fwd(
        &self,
        x: &GpuTensor,
        stream: &Arc<CudaStream>,
    ) -> Result<(GpuTensor, GpuTensor), MLError> {
        let n = x.numel();
        let y = GpuTensor::zeros(&x.shape, stream)?;
        let saved_x = clone_gpu_tensor(x, stream)?;
        let n_i32 = n as i32;
        let cfg = elem_cfg(n);
        unsafe {
            stream
                .launch_builder(&self.gelu_forward)
                .arg(&x.data)
                .arg(&y.data)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("gelu_forward: {e}")))?;
        }
        Ok((y, saved_x))
    }

    /// GELU backward.
    pub fn gelu_bwd(
        &self,
        dy: &GpuTensor,
        saved_x: &GpuTensor,
        stream: &Arc<CudaStream>,
    ) -> Result<GpuTensor, MLError> {
        let n = dy.numel();
        let dx = GpuTensor::zeros(&dy.shape, stream)?;
        let n_i32 = n as i32;
        let cfg = elem_cfg(n);
        unsafe {
            stream
                .launch_builder(&self.gelu_backward)
                .arg(&dy.data)
                .arg(&saved_x.data)
                .arg(&dx.data)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("gelu_backward: {e}")))?;
        }
        Ok(dx)
    }

    // ── Sigmoid ──────────────────────────────────────────────────────

    /// Sigmoid forward: `y = 1 / (1 + exp(-x))`. Saves output for backward.
    pub fn sigmoid_fwd(
        &self,
        x: &GpuTensor,
        stream: &Arc<CudaStream>,
    ) -> Result<(GpuTensor, GpuTensor), MLError> {
        let n = x.numel();
        let y = GpuTensor::zeros(&x.shape, stream)?;
        let n_i32 = n as i32;
        let cfg = elem_cfg(n);
        unsafe {
            stream
                .launch_builder(&self.sigmoid_forward)
                .arg(&x.data)
                .arg(&y.data)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("sigmoid_forward: {e}")))?;
        }
        // Sigmoid backward uses output y, so clone it as saved
        let saved_y = clone_gpu_tensor(&y, stream)?;
        Ok((y, saved_y))
    }

    /// Sigmoid backward: `dx = dy * y * (1 - y)`.
    pub fn sigmoid_bwd(
        &self,
        dy: &GpuTensor,
        saved_y: &GpuTensor,
        stream: &Arc<CudaStream>,
    ) -> Result<GpuTensor, MLError> {
        let n = dy.numel();
        let dx = GpuTensor::zeros(&dy.shape, stream)?;
        let n_i32 = n as i32;
        let cfg = elem_cfg(n);
        unsafe {
            stream
                .launch_builder(&self.sigmoid_backward)
                .arg(&dy.data)
                .arg(&saved_y.data)
                .arg(&dx.data)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("sigmoid_backward: {e}")))?;
        }
        Ok(dx)
    }

    // ── Tanh ─────────────────────────────────────────────────────────

    /// Tanh forward. Saves output for backward.
    pub fn tanh_fwd(
        &self,
        x: &GpuTensor,
        stream: &Arc<CudaStream>,
    ) -> Result<(GpuTensor, GpuTensor), MLError> {
        let n = x.numel();
        let y = GpuTensor::zeros(&x.shape, stream)?;
        let n_i32 = n as i32;
        let cfg = elem_cfg(n);
        unsafe {
            stream
                .launch_builder(&self.tanh_forward)
                .arg(&x.data)
                .arg(&y.data)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("tanh_forward: {e}")))?;
        }
        let saved_y = clone_gpu_tensor(&y, stream)?;
        Ok((y, saved_y))
    }

    /// Tanh backward: `dx = dy * (1 - y^2)`.
    pub fn tanh_bwd(
        &self,
        dy: &GpuTensor,
        saved_y: &GpuTensor,
        stream: &Arc<CudaStream>,
    ) -> Result<GpuTensor, MLError> {
        let n = dy.numel();
        let dx = GpuTensor::zeros(&dy.shape, stream)?;
        let n_i32 = n as i32;
        let cfg = elem_cfg(n);
        unsafe {
            stream
                .launch_builder(&self.tanh_backward)
                .arg(&dy.data)
                .arg(&saved_y.data)
                .arg(&dx.data)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("tanh_backward: {e}")))?;
        }
        Ok(dx)
    }

    // ── Raw CudaSlice methods (for ml-ppo CudaVec interop) ───────────

    /// ReLU forward on raw `CudaSlice`. Returns output slice only (no mask).
    ///
    /// Use this when the backward pass is not needed (inference / PPO forward).
    pub fn relu_fwd_raw(
        &self,
        x: &CudaSlice<f32>,
        len: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<CudaSlice<f32>, MLError> {
        let y = stream.alloc_zeros::<f32>(len).map_err(|e| {
            MLError::ModelError(format!("relu_fwd_raw alloc: {e}"))
        })?;
        let mask = stream.alloc_zeros::<f32>(len).map_err(|e| {
            MLError::ModelError(format!("relu_fwd_raw mask alloc: {e}"))
        })?;
        let n_i32 = len as i32;
        let cfg = elem_cfg(len);
        unsafe {
            stream
                .launch_builder(&self.relu_forward)
                .arg(x)
                .arg(&y)
                .arg(&mask)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("relu_fwd_raw: {e}")))?;
        }
        Ok(y)
    }

    /// ReLU backward on raw `CudaSlice` pairs. Returns `grad_input`.
    pub fn relu_bwd_raw(
        &self,
        grad_output: &CudaSlice<f32>,
        input: &CudaSlice<f32>,
        len: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<CudaSlice<f32>, MLError> {
        // Forward to get the mask
        let mask = stream.alloc_zeros::<f32>(len).map_err(|e| {
            MLError::ModelError(format!("relu_bwd_raw mask alloc: {e}"))
        })?;
        let y_scratch = stream.alloc_zeros::<f32>(len).map_err(|e| {
            MLError::ModelError(format!("relu_bwd_raw y_scratch alloc: {e}"))
        })?;
        let n_i32 = len as i32;
        let cfg = elem_cfg(len);
        unsafe {
            stream
                .launch_builder(&self.relu_forward)
                .arg(input)
                .arg(&y_scratch)
                .arg(&mask)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("relu_bwd_raw fwd: {e}")))?;
        }
        // Backward using the mask
        let dx = stream.alloc_zeros::<f32>(len).map_err(|e| {
            MLError::ModelError(format!("relu_bwd_raw dx alloc: {e}"))
        })?;
        unsafe {
            stream
                .launch_builder(&self.relu_backward)
                .arg(grad_output)
                .arg(&mask)
                .arg(&dx)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("relu_bwd_raw: {e}")))?;
        }
        Ok(dx)
    }

    /// Tanh forward on raw `CudaSlice`. Returns output only (no saved tensor).
    pub fn tanh_fwd_raw(
        &self,
        x: &CudaSlice<f32>,
        len: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<CudaSlice<f32>, MLError> {
        let y = stream.alloc_zeros::<f32>(len).map_err(|e| {
            MLError::ModelError(format!("tanh_fwd_raw alloc: {e}"))
        })?;
        let n_i32 = len as i32;
        let cfg = elem_cfg(len);
        unsafe {
            stream
                .launch_builder(&self.tanh_forward)
                .arg(x)
                .arg(&y)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("tanh_fwd_raw: {e}")))?;
        }
        Ok(y)
    }

    /// Sigmoid forward on raw `CudaSlice`. Returns output only (no saved tensor).
    pub fn sigmoid_fwd_raw(
        &self,
        x: &CudaSlice<f32>,
        len: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<CudaSlice<f32>, MLError> {
        let y = stream.alloc_zeros::<f32>(len).map_err(|e| {
            MLError::ModelError(format!("sigmoid_fwd_raw alloc: {e}"))
        })?;
        let n_i32 = len as i32;
        let cfg = elem_cfg(len);
        unsafe {
            stream
                .launch_builder(&self.sigmoid_forward)
                .arg(x)
                .arg(&y)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("sigmoid_fwd_raw: {e}")))?;
        }
        Ok(y)
    }
}

// ── CUDA source for all activation kernels ───────────────────────────────

const ACTIVATION_CUDA_SRC: &str = r#"
// ═══════════════════════════════════════════════════════════════════════
// Activation kernels -- forward and backward for ReLU, LeakyReLU, GELU,
// Sigmoid, Tanh. Each pair uses the same element-wise threading pattern.
// ═══════════════════════════════════════════════════════════════════════

extern "C" __global__
void relu_forward(const float* __restrict__ x,
                  float* __restrict__ y,
                  float* __restrict__ mask,
                  int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float xi = x[i];
        float m = (xi > 0.0f) ? 1.0f : 0.0f;
        y[i] = xi * m;
        mask[i] = m;
    }
}

extern "C" __global__
void relu_backward(const float* __restrict__ dy,
                   const float* __restrict__ mask,
                   float* __restrict__ dx,
                   int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        dx[i] = dy[i] * mask[i];
    }
}

extern "C" __global__
void leaky_relu_forward(const float* __restrict__ x,
                        float* __restrict__ y,
                        float* __restrict__ mask,
                        float alpha,
                        int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float xi = x[i];
        float m = (xi > 0.0f) ? 1.0f : 0.0f;
        y[i] = (xi > 0.0f) ? xi : alpha * xi;
        mask[i] = m;
    }
}

extern "C" __global__
void leaky_relu_backward(const float* __restrict__ dy,
                         const float* __restrict__ mask,
                         float* __restrict__ dx,
                         float alpha,
                         int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float m = mask[i];
        dx[i] = dy[i] * (m + alpha * (1.0f - m));
    }
}

extern "C" __global__
void gelu_forward(const float* __restrict__ x,
                  float* __restrict__ y,
                  int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float xi = x[i];
        // Approximate GELU: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
        float c = 0.7978845608f; // sqrt(2/pi)
        float inner = c * (xi + 0.044715f * xi * xi * xi);
        y[i] = 0.5f * xi * (1.0f + tanhf(inner));
    }
}

extern "C" __global__
void gelu_backward(const float* __restrict__ dy,
                   const float* __restrict__ x,
                   float* __restrict__ dx,
                   int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float xi = x[i];
        float c = 0.7978845608f;
        float x3 = xi * xi * xi;
        float inner = c * (xi + 0.044715f * x3);
        float tanh_val = tanhf(inner);
        float sech2 = 1.0f - tanh_val * tanh_val;
        float d_inner = c * (1.0f + 3.0f * 0.044715f * xi * xi);
        // d/dx GELU = 0.5 * (1 + tanh(inner)) + 0.5 * x * sech^2(inner) * d_inner
        float grad = 0.5f * (1.0f + tanh_val) + 0.5f * xi * sech2 * d_inner;
        dx[i] = dy[i] * grad;
    }
}

extern "C" __global__
void sigmoid_forward(const float* __restrict__ x,
                     float* __restrict__ y,
                     int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        y[i] = 1.0f / (1.0f + expf(-x[i]));
    }
}

extern "C" __global__
void sigmoid_backward(const float* __restrict__ dy,
                      const float* __restrict__ y,
                      float* __restrict__ dx,
                      int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float yi = y[i];
        dx[i] = dy[i] * yi * (1.0f - yi);
    }
}

extern "C" __global__
void tanh_forward(const float* __restrict__ x,
                  float* __restrict__ y,
                  int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        y[i] = tanhf(x[i]);
    }
}

extern "C" __global__
void tanh_backward(const float* __restrict__ dy,
                   const float* __restrict__ y,
                   float* __restrict__ dx,
                   int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float yi = y[i];
        dx[i] = dy[i] * (1.0f - yi * yi);
    }
}
"#;

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

    fn make_stream() -> Arc<CudaStream> {
        cudarc::driver::CudaContext::new(0)
                .expect("CUDA required")
                .new_stream()
                .expect("fork stream")
    }

    #[test]
    fn test_relu_forward_backward() {
        let stream = make_stream();
        let kernels = ActivationKernels::new(&stream).unwrap();

        let input = vec![-2.0_f32, -1.0, 0.0, 1.0, 2.0];
        let x = GpuTensor::from_host(&input, vec![5], &stream).unwrap();

        let (y, mask) = kernels.relu_fwd(&x, &stream).unwrap();
        let y_host = y.to_host(&stream).unwrap();
        assert_eq!(y_host, vec![0.0, 0.0, 0.0, 1.0, 2.0]);

        let dy = GpuTensor::from_host(&vec![1.0_f32; 5], vec![5], &stream).unwrap();
        let dx = kernels.relu_bwd(&dy, &mask, &stream).unwrap();
        let dx_host = dx.to_host(&stream).unwrap();
        assert_eq!(dx_host, vec![0.0, 0.0, 0.0, 1.0, 1.0]);
    }

    #[test]
    fn test_sigmoid_forward() {
        let stream = make_stream();
        let kernels = ActivationKernels::new(&stream).unwrap();

        let input = vec![0.0_f32];
        let x = GpuTensor::from_host(&input, vec![1], &stream).unwrap();
        let (y, _) = kernels.sigmoid_fwd(&x, &stream).unwrap();
        let y_host = y.to_host(&stream).unwrap();
        assert!((y_host[0] - 0.5).abs() < 1e-5, "sigmoid(0) should be 0.5, got {}", y_host[0]);
    }

    #[test]
    fn test_tanh_forward() {
        let stream = make_stream();
        let kernels = ActivationKernels::new(&stream).unwrap();

        let input = vec![0.0_f32, 1.0];
        let x = GpuTensor::from_host(&input, vec![2], &stream).unwrap();
        let (y, _) = kernels.tanh_fwd(&x, &stream).unwrap();
        let y_host = y.to_host(&stream).unwrap();
        assert!((y_host[0] - 0.0).abs() < 1e-5);
        assert!((y_host[1] - 0.7615942).abs() < 1e-4);
    }
}
