//! CUDA activation function wrappers (f32 native, precompiled cubin).

use std::sync::Arc;

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

use ml_core::MLError;

use super::tensor_util::CudaVec;

struct ActivationKernelSet {
    relu_fwd: cudarc::driver::CudaFunction,
    relu_bwd: cudarc::driver::CudaFunction,
    tanh_fwd: cudarc::driver::CudaFunction,
    sigmoid_fwd: cudarc::driver::CudaFunction,
}

fn load_activation_kernels(stream: &Arc<CudaStream>) -> Result<ActivationKernelSet, MLError> {
    static CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/activation_kernels.cubin"));
    let module = stream.context()
        .load_cubin(CUBIN.to_vec())
        .map_err(|e| MLError::InitializationError {
            component: "activations".to_owned(),
            message: format!("Failed to load activation cubin: {e}"),
        })?;
    let load = |name: &str| -> Result<cudarc::driver::CudaFunction, MLError> {
        module.load_function(name).map_err(|e| MLError::InitializationError {
            component: "activations".to_owned(),
            message: format!("Failed to load {name}: {e}"),
        })
    };
    Ok(ActivationKernelSet {
        relu_fwd: load("relu_forward")?,
        relu_bwd: load("relu_backward")?,
        tanh_fwd: load("tanh_forward")?,
        sigmoid_fwd: load("sigmoid_forward")?,
    })
}

fn launch_cfg(n: usize) -> LaunchConfig {
    let blocks = ((n + 255) / 256) as u32;
    LaunchConfig {
        grid_dim: (blocks, 1, 1),
        block_dim: (256, 1, 1),
        shared_mem_bytes: 0,
    }
}

/// Apply `ReLU` activation: output = max(0, input)
pub fn cuda_relu(
    stream: &Arc<CudaStream>,
    input: &CudaSlice<f32>,
    len: usize,
) -> Result<CudaVec, MLError> {
    let kernels = load_activation_kernels(stream)?;
    let mut output = stream.alloc_zeros::<f32>(len).map_err(|e| {
        MLError::ModelError(format!("alloc relu output: {e}"))
    })?;
    let n = len as i32;
    unsafe {
        stream.launch_builder(&kernels.relu_fwd)
            .arg(&mut output)
            .arg(input)
            .arg(&n)
            .launch(launch_cfg(len))
            .map_err(|e| MLError::ModelError(format!("relu_forward launch: {e}")))?;
    }
    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 = load_activation_kernels(stream)?;
    let mut grad_input = stream.alloc_zeros::<f32>(len).map_err(|e| {
        MLError::ModelError(format!("alloc relu_backward output: {e}"))
    })?;
    let n = len as i32;
    unsafe {
        stream.launch_builder(&kernels.relu_bwd)
            .arg(&mut grad_input)
            .arg(grad_output)
            .arg(input)
            .arg(&n)
            .launch(launch_cfg(len))
            .map_err(|e| MLError::ModelError(format!("relu_backward launch: {e}")))?;
    }
    Ok(CudaVec::new(grad_input, len))
}

/// Apply tanh activation: output = tanh(input)
pub fn cuda_tanh(
    stream: &Arc<CudaStream>,
    input: &CudaSlice<f32>,
    len: usize,
) -> Result<CudaVec, MLError> {
    let kernels = load_activation_kernels(stream)?;
    let mut output = stream.alloc_zeros::<f32>(len).map_err(|e| {
        MLError::ModelError(format!("alloc tanh output: {e}"))
    })?;
    let n = len as i32;
    unsafe {
        stream.launch_builder(&kernels.tanh_fwd)
            .arg(&mut output)
            .arg(input)
            .arg(&n)
            .launch(launch_cfg(len))
            .map_err(|e| MLError::ModelError(format!("tanh_forward launch: {e}")))?;
    }
    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 = load_activation_kernels(stream)?;
    let mut output = stream.alloc_zeros::<f32>(len).map_err(|e| {
        MLError::ModelError(format!("alloc sigmoid output: {e}"))
    })?;
    let n = len as i32;
    unsafe {
        stream.launch_builder(&kernels.sigmoid_fwd)
            .arg(&mut output)
            .arg(input)
            .arg(&n)
            .launch(launch_cfg(len))
            .map_err(|e| MLError::ModelError(format!("sigmoid_forward launch: {e}")))?;
    }
    Ok(CudaVec::new(output, len))
}
