//! GPU LayerNorm implementation using CUDA kernels.
//!
//! Computes per-sample mean and variance, normalizes, then applies learned
//! scale (gamma) and shift (beta) parameters stored in a `GpuVarStore`.
//!
//! This replaces `candle_nn::LayerNorm` with a pure cudarc implementation.

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::var_store::GpuVarStore;
use super::init;

/// GPU LayerNorm layer.
///
/// Normalizes over the last dimension (features) of a 2D `[batch, features]` tensor.
/// Stores weight (gamma) and bias (beta) parameters in a `GpuVarStore`.
pub struct GpuLayerNorm {
    /// Name prefix for parameters in the var store.
    pub weight_name: String,
    pub bias_name: String,
    /// Number of features (normalized dimension).
    pub features: usize,
    /// Epsilon for numerical stability.
    pub eps: f32,
    /// Compiled forward kernel.
    kernel: Option<CudaFunction>,
}

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

impl GpuLayerNorm {
    /// Register a new LayerNorm in the given var store.
    ///
    /// Initializes weight (gamma) to ones and bias (beta) to zeros.
    pub fn new(
        prefix: &str,
        features: usize,
        eps: f32,
        store: &mut GpuVarStore,
    ) -> Result<Self, MLError> {
        let weight_name = format!("{prefix}.weight");
        let bias_name = format!("{prefix}.bias");

        // gamma = ones
        let ones = vec![1.0_f32; features];
        let w_data = init::upload_to_gpu(&ones, store.stream())?;
        store.register(weight_name.clone(), w_data, vec![features])?;

        // beta = zeros
        let b_data = init::zeros(features, store.stream())?;
        store.register(bias_name.clone(), b_data, vec![features])?;

        Ok(Self {
            weight_name,
            bias_name,
            features,
            eps,
            kernel: None,
        })
    }

    /// Forward pass: LayerNorm over the last dimension.
    ///
    /// Input shape: `[batch, features]`.
    /// Output shape: `[batch, features]`.
    pub fn forward(
        &mut self,
        x: &GpuTensor,
        store: &GpuVarStore,
        stream: &Arc<CudaStream>,
    ) -> Result<GpuTensor, MLError> {
        let batch = if x.ndim() == 1 { 1 } else { x.shape()[0] };
        let features = self.features;

        if x.numel() != batch * features {
            return Err(MLError::DimensionMismatch {
                expected: batch * features,
                actual: x.numel(),
            });
        }

        let w_param = store.get(&self.weight_name).ok_or_else(|| {
            MLError::ModelError(format!("LayerNorm weight '{}' not found", self.weight_name))
        })?;
        let b_param = store.get(&self.bias_name).ok_or_else(|| {
            MLError::ModelError(format!("LayerNorm bias '{}' not found", self.bias_name))
        })?;

        // Compile kernel on first use
        if self.kernel.is_none() {
            self.compile_kernel(stream)?;
        }

        let y = GpuTensor::zeros(&x.shape, stream)?;
        let batch_i32 = batch as i32;
        let features_i32 = features as i32;

        // Each block handles one sample (row). Use enough threads to cover features.
        // For simplicity, cap threads at 256 and have each thread handle multiple features.
        let threads = 256_u32.min(features as u32);
        let cfg = LaunchConfig {
            grid_dim: (batch as u32, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: (threads as usize * std::mem::size_of::<f32>() * 2) as u32,
        };

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

        unsafe {
            stream
                .launch_builder(kernel)
                .arg(&x.data)
                .arg(&y.data)
                .arg(&w_param.data)
                .arg(&b_param.data)
                .arg(&batch_i32)
                .arg(&features_i32)
                .arg(&self.eps)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("layer_norm_forward: {e}")))?;
        }

        Ok(y)
    }

    /// Compile the LayerNorm 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(LAYER_NORM_CUDA_SRC, &context)
            .map_err(|e| MLError::ModelError(format!("layer_norm kernel compilation: {e}")))?;
        let module = context.load_module(ptx)
            .map_err(|e| MLError::ModelError(format!("layer_norm module load: {e}")))?;
        let kernel = module.load_function("layer_norm_forward")
            .map_err(|e| MLError::ModelError(format!("layer_norm_forward load: {e}")))?;
        self.kernel = Some(kernel);
        Ok(())
    }
}

// ── CUDA source for LayerNorm kernel ────────────────────────────────────

const LAYER_NORM_CUDA_SRC: &str = r#"
// LayerNorm forward kernel.
// Each block processes one sample (row) of the [batch, features] input.
// Two-pass approach:
//   Pass 1: compute mean and variance via parallel reduction in shared memory.
//   Pass 2: normalize, scale, and shift each element.

extern "C" __global__
void layer_norm_forward(const float* __restrict__ x,
                        float* __restrict__ y,
                        const float* __restrict__ gamma,
                        const float* __restrict__ beta,
                        int batch,
                        int features,
                        float eps) {
    int row = blockIdx.x;
    if (row >= batch) return;

    const float* x_row = x + row * features;
    float* y_row = y + row * features;

    // Shared memory: first half for sum, second half for sum of squares
    extern __shared__ float smem[];
    float* s_sum = smem;
    float* s_sq  = smem + blockDim.x;

    // Thread-local accumulation over features (each thread handles multiple features)
    float local_sum = 0.0f;
    float local_sq  = 0.0f;
    for (int j = threadIdx.x; j < features; j += blockDim.x) {
        float val = x_row[j];
        local_sum += val;
        local_sq  += val * val;
    }

    s_sum[threadIdx.x] = local_sum;
    s_sq[threadIdx.x]  = local_sq;
    __syncthreads();

    // Parallel reduction
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            s_sum[threadIdx.x] += s_sum[threadIdx.x + stride];
            s_sq[threadIdx.x]  += s_sq[threadIdx.x + stride];
        }
        __syncthreads();
    }

    float mean = s_sum[0] / (float)features;
    float var  = s_sq[0] / (float)features - mean * mean;
    float inv_std = rsqrtf(var + eps);

    // Normalize, scale, shift
    for (int j = threadIdx.x; j < features; j += blockDim.x) {
        float normalized = (x_row[j] - mean) * inv_std;
        y_row[j] = normalized * gamma[j] + beta[j];
    }
}
"#;
