//! CUDA softmax and log-softmax kernels.
//!
//! Numerically stable implementations operating on row-major `[batch, dim]` data.

use std::sync::{Arc, OnceLock};

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

use ml_core::MLError;

use super::tensor_util::CudaVec;

/// Cached PTX for softmax kernels -- compiled once per process via `OnceLock`.
/// NVRTC compilation is ~50ms per call; caching eliminates this from every
/// forward pass (`cuda_softmax` and `cuda_log_softmax` are called per-batch).
static SOFTMAX_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();

/// Softmax and log-softmax CUDA kernels.
///
/// Each warp/block handles one row of [dim] elements.
/// Uses the numerically stable formulation: `softmax(x_i)` = `exp(x_i - max(x))` / `sum(exp(x - max(x)))`
const SOFTMAX_KERNEL: &str = r#"
extern "C" __global__ void softmax_forward(
    float* __restrict__ output,
    const float* __restrict__ input,
    int batch_size,
    int dim
) {
    int row = blockIdx.x;
    if (row >= batch_size) return;

    const float* in_row = input + row * dim;
    float* out_row = output + row * dim;

    // Find max for numerical stability
    float max_val = -1e30f;
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        float v = in_row[i];
        if (v > max_val) max_val = v;
    }
    // Warp reduction for max
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        float other = __shfl_down_sync(0xffffffff, max_val, offset);
        if (other > max_val) max_val = other;
    }
    // Broadcast max from lane 0
    max_val = __shfl_sync(0xffffffff, max_val, 0);

    // Compute exp(x - max) and sum
    float sum = 0.0f;
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        float e = expf(in_row[i] - max_val);
        out_row[i] = e;
        sum += e;
    }
    // Warp reduction for sum
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(0xffffffff, sum, offset);
    }
    sum = __shfl_sync(0xffffffff, sum, 0);

    // Normalize
    float inv_sum = 1.0f / sum;
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        out_row[i] *= inv_sum;
    }
}

extern "C" __global__ void log_softmax_forward(
    float* __restrict__ output,
    const float* __restrict__ input,
    int batch_size,
    int dim
) {
    int row = blockIdx.x;
    if (row >= batch_size) return;

    const float* in_row = input + row * dim;
    float* out_row = output + row * dim;

    // Find max for numerical stability
    float max_val = -1e30f;
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        float v = in_row[i];
        if (v > max_val) max_val = v;
    }
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        float other = __shfl_down_sync(0xffffffff, max_val, offset);
        if (other > max_val) max_val = other;
    }
    max_val = __shfl_sync(0xffffffff, max_val, 0);

    // Compute sum of exp(x - max)
    float sum = 0.0f;
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        sum += expf(in_row[i] - max_val);
    }
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(0xffffffff, sum, offset);
    }
    sum = __shfl_sync(0xffffffff, sum, 0);

    float log_sum = logf(sum);

    // log_softmax(x_i) = x_i - max - log(sum(exp(x - max)))
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        out_row[i] = in_row[i] - max_val - log_sum;
    }
}
"#;

struct SoftmaxKernels {
    softmax_fwd: cudarc::driver::CudaFunction,
    log_softmax_fwd: cudarc::driver::CudaFunction,
}

fn compile_softmax_kernels(stream: &Arc<CudaStream>) -> Result<SoftmaxKernels, MLError> {
    let context = stream.context();
    let ptx_result = SOFTMAX_PTX.get_or_init(|| {
        ml_core::cuda_compile::compile_ptx_for_device(SOFTMAX_KERNEL, context)
            .map_err(|e| format!("Failed to compile softmax kernels: {e}"))
    });
    let ptx = ptx_result.as_ref().map_err(|e| {
        MLError::InitializationError {
            component: "softmax".to_owned(),
            message: e.clone(),
        }
    })?;
    let module = context.load_module(ptx.clone()).map_err(|e| {
        MLError::InitializationError {
            component: "softmax".to_owned(),
            message: format!("Failed to load softmax module: {e}"),
        }
    })?;
    Ok(SoftmaxKernels {
        softmax_fwd: module.load_function("softmax_forward").map_err(|e| {
            MLError::InitializationError {
                component: "softmax".to_owned(),
                message: format!("Failed to load softmax_forward: {e}"),
            }
        })?,
        log_softmax_fwd: module.load_function("log_softmax_forward").map_err(|e| {
            MLError::InitializationError {
                component: "softmax".to_owned(),
                message: format!("Failed to load log_softmax_forward: {e}"),
            }
        })?,
    })
}

/// Compute softmax over the last dimension: `[batch_size, dim]`
pub fn cuda_softmax(
    stream: &Arc<CudaStream>,
    input: &CudaSlice<f32>,
    batch_size: usize,
    dim: usize,
) -> Result<CudaVec, MLError> {
    let kernels = compile_softmax_kernels(stream)?;
    let total = batch_size * dim;
    let mut output = stream.alloc_zeros::<f32>(total).map_err(|e| {
        MLError::ModelError(format!("Failed to alloc softmax output: {e}"))
    })?;

    // One block per row, warp-width threads (32 for dims <= 32, else 32)
    let threads = dim.min(32) as u32;
    let batch_i32 = batch_size as i32;
    let dim_i32 = dim as i32;

    // SAFETY: softmax_forward reads batch_size*dim f32 input elements and
    // writes batch_size*dim output elements. One block processes each row.
    // Both buffers are pre-allocated with `total` = batch_size*dim elements.
    unsafe {
        stream.launch_builder(&kernels.softmax_fwd)
            .arg(&mut output)
            .arg(input)
            .arg(&batch_i32)
            .arg(&dim_i32)
            .launch(LaunchConfig {
                grid_dim: (batch_size as u32, 1, 1),
                block_dim: (threads, 1, 1),
                shared_mem_bytes: 0,
            })
            .map_err(|e| MLError::ModelError(format!("softmax_forward launch failed: {e}")))?;
    }

    Ok(CudaVec::new(output, total))
}

/// Compute log-softmax over the last dimension: `[batch_size, dim]`
pub fn cuda_log_softmax(
    stream: &Arc<CudaStream>,
    input: &CudaSlice<f32>,
    batch_size: usize,
    dim: usize,
) -> Result<CudaVec, MLError> {
    let kernels = compile_softmax_kernels(stream)?;
    let total = batch_size * dim;
    let mut output = stream.alloc_zeros::<f32>(total).map_err(|e| {
        MLError::ModelError(format!("Failed to alloc log_softmax output: {e}"))
    })?;

    let threads = dim.min(32) as u32;
    let batch_i32 = batch_size as i32;
    let dim_i32 = dim as i32;

    // SAFETY: log_softmax_forward reads batch_size*dim f32 input elements and
    // writes batch_size*dim output elements. One block per row, same invariants
    // as softmax_forward above.
    unsafe {
        stream.launch_builder(&kernels.log_softmax_fwd)
            .arg(&mut output)
            .arg(input)
            .arg(&batch_i32)
            .arg(&dim_i32)
            .launch(LaunchConfig {
                grid_dim: (batch_size as u32, 1, 1),
                block_dim: (threads, 1, 1),
                shared_mem_bytes: 0,
            })
            .map_err(|e| MLError::ModelError(format!("log_softmax_forward launch failed: {e}")))?;
    }

    Ok(CudaVec::new(output, total))
}
