#![allow(unsafe_code)] // Required for CUDA kernel launches via cudarc.

//! `RMSNorm` - Root Mean Square Layer Normalization
//!
//! `RMSNorm` is a faster alternative to `LayerNorm` that achieves ~15% speedup
//! with similar performance. It normalizes by the root mean square instead
//! of mean and variance, reducing computation.
//!
//! Paper: "Root Mean Square Layer Normalization" (Zhang & Sennrich, 2019)
//! Wave 26 P2.4
//!
//! ## Weight storage
//!
//! Learnable parameters are stored in a [`GpuVarStore`] (native CUDA) instead
//! of Candle's `VarMap`. Forward passes run entirely on GPU via compiled CUDA
//! kernels -- zero host downloads.

use std::sync::Arc;

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

use ml_core::cuda_autograd::{GpuTensor, GpuVarStore};
use ml_core::MLError;

/// Normalization type configuration
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NormType {
    /// Standard `LayerNorm` (mean and variance normalization)
    LayerNorm,
    /// `RMSNorm` (root mean square normalization, ~15% faster)
    RMSNorm,
    /// No normalization
    None,
}

impl Default for NormType {
    fn default() -> Self {
        Self::LayerNorm
    }
}

// ── RMSNorm CUDA kernel source ──────────────────────────────────────────

const RMSNORM_CUDA_SRC: &str = r#"
// RMSNorm forward kernel.
// Each block processes one sample (row) of the [batch, features] input.
// Computes: y = x / sqrt(mean(x^2) + eps) * weight
//
// Two-pass approach:
//   Pass 1: compute mean(x^2) via parallel reduction in shared memory.
//   Pass 2: normalize and scale each element.

extern "C" __global__
void rmsnorm_forward(const float* __restrict__ x,
                     float* __restrict__ y,
                     const float* __restrict__ weight,
                     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;

    extern __shared__ float smem[];

    // Thread-local accumulation of x^2 over features
    float local_sq = 0.0f;
    for (int j = threadIdx.x; j < features; j += blockDim.x) {
        float val = x_row[j];
        local_sq += val * val;
    }

    smem[threadIdx.x] = local_sq;
    __syncthreads();

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

    // rms = sqrt(mean(x^2) + eps)
    float inv_rms = rsqrtf(smem[0] / (float)features + eps);

    // Normalize and scale
    for (int j = threadIdx.x; j < features; j += blockDim.x) {
        y_row[j] = x_row[j] * inv_rms * weight[j];
    }
}
"#;

// ── LayerNorm CUDA kernel source (same as ml-core GpuLayerNorm) ──────────

const LAYERNORM_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 layernorm_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
    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];
    }
}
"#;

/// Compile a CUDA kernel from source and return the function handle.
fn compile_kernel(
    src: &str,
    fn_name: &str,
    stream: &Arc<CudaStream>,
) -> Result<CudaFunction, MLError> {
    let context = stream.context();
    let ptx: Ptx = ml_core::cuda_compile::compile_ptx_for_device(src, context)
        .map_err(|e| MLError::ModelError(format!("{fn_name} kernel compilation: {e}")))?;
    let module = context.load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("{fn_name} module load: {e}")))?;
    module.load_function(fn_name)
        .map_err(|e| MLError::ModelError(format!("{fn_name} load: {e}")))
}

/// `RMSNorm` layer - faster alternative to `LayerNorm`
///
/// Weights are stored in a [`GpuVarStore`] as native CUDA buffers.
/// Forward pass runs entirely on GPU via a compiled CUDA kernel -- zero host
/// downloads.
#[allow(missing_debug_implementations)]
pub struct RMSNorm {
    store: GpuVarStore,
    stream: Arc<CudaStream>,
    eps: f64,
    dim: usize,
    /// Eagerly compiled `RMSNorm` CUDA kernel.
    kernel: CudaFunction,
}

impl RMSNorm {
    /// Create a new `RMSNorm` layer with native CUDA weight storage.
    ///
    /// The CUDA kernel is compiled eagerly so `forward()` can take `&self`.
    pub fn new(stream: Arc<CudaStream>, dim: usize, eps: f64) -> Result<Self, MLError> {
        let mut store = GpuVarStore::new(Arc::clone(&stream));
        let ones = vec![1.0_f32; dim];
        let weight_data = ml_core::cuda_autograd::init::upload_to_gpu(&ones, &stream)?;
        store.register("weight", weight_data, vec![dim])?;
        let kernel = compile_kernel(RMSNORM_CUDA_SRC, "rmsnorm_forward", &stream)?;
        Ok(Self { store, stream, eps, dim, kernel })
    }

    /// Create with default epsilon (1e-6).
    pub fn new_default(stream: Arc<CudaStream>, dim: usize) -> Result<Self, MLError> {
        Self::new(stream, dim, 1e-6)
    }

    /// Forward pass -- GPU-only via compiled CUDA kernel.
    ///
    /// Computes: `x / sqrt(mean(x^2) + eps) * weight`
    ///
    /// Input/output shape: `[batch, ..., dim]` where last dimension = `self.dim`.
    pub fn forward(&self, x: &GpuTensor) -> Result<GpuTensor, MLError> {
        let shape = x.shape().to_vec();
        let dim = self.dim;
        let total = x.numel();

        if total == 0 || dim == 0 {
            return x.gpu_clone(&self.stream);
        }

        let batch = total / dim;

        let weight_param = self.store.get("weight").ok_or_else(|| {
            MLError::ModelError("RMSNorm param 'weight' not found".to_owned())
        })?;

        let y = GpuTensor::zeros(&shape, &self.stream)?;
        let batch_i32 = batch as i32;
        let features_i32 = dim as i32;
        let eps_f32 = self.eps as f32;

        let threads = 256_u32.min(dim as u32);
        let cfg = LaunchConfig {
            grid_dim: (batch as u32, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: threads * std::mem::size_of::<f32>() as u32,
        };

        // SAFETY: x, y, weight_param.data are valid device allocations. Kernel indices bounds-checked.
        unsafe {
            self.stream
                .launch_builder(&self.kernel)
                .arg(x.data())
                .arg(y.data())
                .arg(&weight_param.data)
                .arg(&batch_i32)
                .arg(&features_i32)
                .arg(&eps_f32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("rmsnorm_forward: {e}")))?;
        }

        Ok(y)
    }

    /// Get the dimension being normalized
    pub const fn dim(&self) -> usize { self.dim }
    /// Get epsilon value
    pub const fn eps(&self) -> f64 { self.eps }
    /// Reference to the underlying `GpuVarStore`.
    pub const fn store(&self) -> &GpuVarStore { &self.store }
    /// Mutable reference to the underlying `GpuVarStore`.
    pub const fn store_mut(&mut self) -> &mut GpuVarStore { &mut self.store }
}

/// `LayerNorm` with native CUDA weight storage.
///
/// Forward pass runs entirely on GPU via a compiled CUDA kernel -- zero host
/// downloads. The kernel is compiled eagerly during construction.
#[allow(missing_debug_implementations)]
pub struct LayerNorm {
    store: GpuVarStore,
    stream: Arc<CudaStream>,
    eps: f64,
    dim: usize,
    /// Weight param name in the store.
    weight_name: String,
    /// Bias param name in the store.
    bias_name: String,
    /// Eagerly compiled `LayerNorm` CUDA kernel.
    kernel: CudaFunction,
}

impl LayerNorm {
    /// Create a new `LayerNorm` layer with native CUDA weight storage.
    ///
    /// The CUDA kernel is compiled eagerly so `forward()` can take `&self`.
    pub fn new(stream: Arc<CudaStream>, dim: usize, eps: f64) -> Result<Self, MLError> {
        let mut store = GpuVarStore::new(Arc::clone(&stream));
        let weight_name = "ln.weight".to_owned();
        let bias_name = "ln.bias".to_owned();

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

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

        let kernel = compile_kernel(LAYERNORM_CUDA_SRC, "layernorm_forward", &stream)?;

        Ok(Self { store, stream, eps, dim, weight_name, bias_name, kernel })
    }

    /// Create with default epsilon (1e-6).
    pub fn new_default(stream: Arc<CudaStream>, dim: usize) -> Result<Self, MLError> {
        Self::new(stream, dim, 1e-6)
    }

    /// Forward pass -- GPU-only via compiled CUDA kernel.
    ///
    /// Computes: `(x - mean(x)) / sqrt(var(x) + eps) * weight + bias`
    pub fn forward(&self, x: &GpuTensor) -> Result<GpuTensor, MLError> {
        let shape = x.shape().to_vec();
        let dim = self.dim;
        let total = x.numel();

        if total == 0 || dim == 0 {
            return x.gpu_clone(&self.stream);
        }

        let batch = total / dim;

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

        let y = GpuTensor::zeros(&shape, &self.stream)?;
        let batch_i32 = batch as i32;
        let features_i32 = dim as i32;
        let eps_f32 = self.eps as f32;

        let threads = 256_u32.min(dim 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,
        };

        // SAFETY: x, y, w_param, b_param are valid device allocations. Kernel indices bounds-checked.
        unsafe {
            self.stream
                .launch_builder(&self.kernel)
                .arg(x.data())
                .arg(y.data())
                .arg(&w_param.data)
                .arg(&b_param.data)
                .arg(&batch_i32)
                .arg(&features_i32)
                .arg(&eps_f32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("layernorm_forward: {e}")))?;
        }

        Ok(y)
    }

    /// Get the dimension being normalized
    pub const fn dim(&self) -> usize { self.dim }
    /// Get epsilon value
    pub const fn eps(&self) -> f64 { self.eps }
    /// Reference to the underlying `GpuVarStore`.
    pub const fn store(&self) -> &GpuVarStore { &self.store }
}

#[cfg(test)]
#[allow(clippy::use_debug, clippy::unnecessary_wraps, clippy::useless_vec)]
mod tests {
    use super::*;
    use std::time::Instant;
    use tracing::info;

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

    #[test]
    fn test_rmsnorm_creation() -> anyhow::Result<()> {
        let rmsnorm = RMSNorm::new_default(make_stream(), 128)?;
        assert_eq!(rmsnorm.dim(), 128);
        assert_eq!(rmsnorm.eps(), 1e-6);
        Ok(())
    }

    #[test]
    fn test_layernorm_creation() -> anyhow::Result<()> {
        let layernorm = LayerNorm::new_default(make_stream(), 128)?;
        assert_eq!(layernorm.dim(), 128);
        assert_eq!(layernorm.eps(), 1e-6);
        Ok(())
    }

    #[test]
    fn test_rmsnorm_forward() -> anyhow::Result<()> {
        let stream = make_stream();
        let (batch_size, dim) = (4, 128);
        let rmsnorm = RMSNorm::new_default(Arc::clone(&stream), dim)?;
        let input_data: Vec<f32> = (0..batch_size * dim).map(|i| (i as f32 * 0.01).sin()).collect();
        let input = GpuTensor::from_host(&input_data, vec![batch_size, dim], &stream)?;
        let output = rmsnorm.forward(&input)?;
        assert_eq!(output.shape(), &[batch_size, dim]);
        let output_host = output.to_host(&stream)?;
        for batch in 0..batch_size {
            let start = batch * dim;
            let row_sq_mean: f64 = output_host[start..start + dim]
                .iter()
                .map(|&v| (v as f64) * (v as f64))
                .sum::<f64>() / dim as f64;
            assert!(row_sq_mean > 0.8 && row_sq_mean < 1.2,
                "RMSNorm output row {batch} should have ~unit variance, got {row_sq_mean}");
        }
        Ok(())
    }

    #[test]
    fn test_layernorm_forward() -> anyhow::Result<()> {
        let stream = make_stream();
        let (batch_size, dim) = (4, 128);
        let layernorm = LayerNorm::new_default(Arc::clone(&stream), dim)?;
        let input_data: Vec<f32> = (0..batch_size * dim).map(|i| (i as f32 * 0.01).sin()).collect();
        let input = GpuTensor::from_host(&input_data, vec![batch_size, dim], &stream)?;
        let output = layernorm.forward(&input)?;
        assert_eq!(output.shape(), &[batch_size, dim]);
        let output_host = output.to_host(&stream)?;
        for batch in 0..batch_size {
            let start = batch * dim;
            let row = &output_host[start..start + dim];
            let mean: f64 = row.iter().map(|&v| v as f64).sum::<f64>() / dim as f64;
            let sq_mean: f64 = row.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>() / dim as f64;
            assert!(mean.abs() < 0.1, "LayerNorm output row {batch} mean should be ~0, got {mean}");
            assert!(sq_mean > 0.8 && sq_mean < 1.2,
                "LayerNorm output row {batch} should have ~unit variance, got {sq_mean}");
        }
        Ok(())
    }

    #[test]
    fn test_rmsnorm_vs_layernorm_performance() -> anyhow::Result<()> {
        let stream = make_stream();
        let (batch_size, dim, iters) = (32, 256, 100);
        let rmsnorm = RMSNorm::new_default(Arc::clone(&stream), dim)?;
        let layernorm = LayerNorm::new_default(Arc::clone(&stream), dim)?;
        let input_data: Vec<f32> = (0..batch_size * dim).map(|i| (i as f32 * 0.01).sin()).collect();
        let input = GpuTensor::from_host(&input_data, vec![batch_size, dim], &stream)?;
        let t0 = Instant::now();
        for _ in 0..iters { let _ = rmsnorm.forward(&input)?; }
        let d0 = t0.elapsed();
        let t1 = Instant::now();
        for _ in 0..iters { let _ = layernorm.forward(&input)?; }
        let d1 = t1.elapsed();
        let speedup = d1.as_secs_f64() / d0.as_secs_f64();
        info!(speedup = %format!("{:.2}x", speedup), "RMSNorm vs LayerNorm");
        assert!(speedup >= 0.20);
        Ok(())
    }

    #[test]
    fn test_rmsnorm_vs_layernorm_numerical_similarity() -> anyhow::Result<()> {
        let stream = make_stream();
        let (batch_size, dim) = (4, 128);
        let rmsnorm = RMSNorm::new_default(Arc::clone(&stream), dim)?;
        let layernorm = LayerNorm::new_default(Arc::clone(&stream), dim)?;
        let input_data: Vec<f32> = (0..batch_size * dim).map(|i| (i as f32 * 0.01).sin()).collect();
        let input = GpuTensor::from_host(&input_data, vec![batch_size, dim], &stream)?;

        let rms_out = rmsnorm.forward(&input)?;
        let ln_out = layernorm.forward(&input)?;
        let rms_host = rms_out.to_host(&stream)?;
        let ln_host = ln_out.to_host(&stream)?;

        // Both should produce values with roughly unit variance
        let rv: f64 = rms_host.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>() / rms_host.len() as f64;
        let lv: f64 = ln_host.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>() / ln_host.len() as f64;
        assert!(rv > 0.5 && rv < 2.0, "RMSNorm overall variance: {rv}");
        assert!(lv > 0.5 && lv < 2.0, "LayerNorm overall variance: {lv}");
        Ok(())
    }

    #[test]
    fn test_norm_type_enum() -> anyhow::Result<()> {
        assert_eq!(NormType::default(), NormType::LayerNorm);
        assert_eq!(vec![NormType::LayerNorm, NormType::RMSNorm, NormType::None].len(), 3);
        Ok(())
    }

    #[test]
    fn test_rmsnorm_3d_input() -> anyhow::Result<()> {
        let stream = make_stream();
        let (bs, sl, dim) = (4, 16, 128);
        let rmsnorm = RMSNorm::new_default(Arc::clone(&stream), dim)?;
        let data: Vec<f32> = (0..bs * sl * dim).map(|i| (i as f32 * 0.01).sin()).collect();
        let input = GpuTensor::from_host(&data, vec![bs, sl, dim], &stream)?;
        assert_eq!(rmsnorm.forward(&input)?.shape(), &[bs, sl, dim]);
        Ok(())
    }

    #[test]
    fn test_layernorm_3d_input() -> anyhow::Result<()> {
        let stream = make_stream();
        let (bs, sl, dim) = (4, 16, 128);
        let layernorm = LayerNorm::new_default(Arc::clone(&stream), dim)?;
        let data: Vec<f32> = (0..bs * sl * dim).map(|i| (i as f32 * 0.01).sin()).collect();
        let input = GpuTensor::from_host(&data, vec![bs, sl, dim], &stream)?;
        assert_eq!(layernorm.forward(&input)?.shape(), &[bs, sl, dim]);
        Ok(())
    }
}
