//! CUDA activation function wrappers.
//!
//! Thin wrappers around `ml_core::cuda_autograd::ActivationKernels` that
//! present the original `CudaVec`-returning API. The underlying CUDA kernels
//! are compiled and executed by ml-core -- no duplicate PTX source here.

use std::sync::Arc;

use cudarc::driver::{CudaSlice, CudaStream};

use ml_core::cuda_autograd::ActivationKernels;
use ml_core::MLError;

use super::tensor_util::CudaVec;

/// Apply `ReLU` activation: output = max(0, input)
pub fn cuda_relu(
    stream: &Arc<CudaStream>,
    input: &CudaSlice<f32>,
    len: usize,
) -> Result<CudaVec, MLError> {
    let kernels = ActivationKernels::new(stream)?;
    let output = kernels.relu_fwd_raw(input, len, stream)?;
    Ok(CudaVec::new(output, len))
}

/// Compute `ReLU` backward: `grad_input` = `grad_output` * (input > 0)
pub fn cuda_relu_backward(
    stream: &Arc<CudaStream>,
    grad_output: &CudaSlice<f32>,
    input: &CudaSlice<f32>,
    len: usize,
) -> Result<CudaVec, MLError> {
    let kernels = ActivationKernels::new(stream)?;
    let dx = kernels.relu_bwd_raw(grad_output, input, len, stream)?;
    Ok(CudaVec::new(dx, len))
}

/// Apply tanh activation: output = tanh(input)
pub fn cuda_tanh(
    stream: &Arc<CudaStream>,
    input: &CudaSlice<f32>,
    len: usize,
) -> Result<CudaVec, MLError> {
    let kernels = ActivationKernels::new(stream)?;
    let output = kernels.tanh_fwd_raw(input, len, stream)?;
    Ok(CudaVec::new(output, len))
}

/// Apply sigmoid activation: output = 1 / (1 + exp(-input))
pub fn cuda_sigmoid(
    stream: &Arc<CudaStream>,
    input: &CudaSlice<f32>,
    len: usize,
) -> Result<CudaVec, MLError> {
    let kernels = ActivationKernels::new(stream)?;
    let output = kernels.sigmoid_fwd_raw(input, len, stream)?;
    Ok(CudaVec::new(output, len))
}
