//! GPU-resident linear (fully connected) layer using cuBLAS sgemm.
//!
//! Replaces `candle_nn::Linear`. Weights and biases are stored as raw
//! `CudaSlice<f32>` with no Candle tensor overhead. Forward pass is a
//! single cuBLAS sgemm + bias-add CUDA kernel.

use std::sync::OnceLock;

use cudarc;
use cudarc::cublas::sys::cublasOperation_t;
use cudarc::driver::{CudaSlice, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;

use ml_core::MLError;

use super::GpuContext;
use super::tensor_util::CudaVec;
use super::{raw_ptr, raw_ptr_mut};

/// Cached PTX for `bias_add` kernel -- compiled once per process via `OnceLock`.
static LINEAR_BIAS_ADD_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();

/// Linear layer: `Y = X @ W^T + b`
///
/// Weight shape: `[out_features, in_features]` (row-major)
/// Bias shape: `[out_features]`
///
/// Forward: for input `X [batch, in_features]`:
///   `Y [batch, out_features] = X @ W^T + broadcast(b)`
pub struct CudaLinear {
    /// Weight matrix `[out_features * in_features]` (row-major)
    pub weight: CudaSlice<f32>,
    /// Bias vector `[out_features]`
    pub bias: CudaSlice<f32>,
    /// Weight gradient `[out_features * in_features]`
    pub weight_grad: CudaSlice<f32>,
    /// Bias gradient `[out_features]`
    pub bias_grad: CudaSlice<f32>,
    /// Input dimension
    pub in_features: usize,
    /// Output dimension
    pub out_features: usize,
    /// Shared CUDA context
    ctx: GpuContext,
    /// Compiled bias-add kernel
    bias_add_func: cudarc::driver::CudaFunction,
}

impl std::fmt::Debug for CudaLinear {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CudaLinear")
            .field("in_features", &self.in_features)
            .field("out_features", &self.out_features)
            .finish()
    }
}

const BIAS_ADD_KERNEL: &str = r#"
extern "C" __global__ void bias_add(
    float* __restrict__ output,
    const float* __restrict__ bias,
    int out_features,
    int total_elements
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total_elements) {
        int col = idx % out_features;
        output[idx] += bias[col];
    }
}
"#;

fn compile_bias_add(ctx: &GpuContext) -> Result<cudarc::driver::CudaFunction, MLError> {
    let context = ctx.stream.context();
    let ptx_result = LINEAR_BIAS_ADD_PTX.get_or_init(|| {
        ml_core::cuda_compile::compile_ptx_for_device(BIAS_ADD_KERNEL, context)
            .map_err(|e| format!("Failed to compile bias_add kernel: {e}"))
    });
    let ptx = ptx_result.as_ref().map_err(|e| {
        MLError::InitializationError {
            component: "CudaLinear".to_owned(),
            message: e.clone(),
        }
    })?;
    let module = context.load_module(ptx.clone()).map_err(|e| {
        MLError::InitializationError {
            component: "CudaLinear".to_owned(),
            message: format!("Failed to load bias_add module: {e}"),
        }
    })?;
    module.load_function("bias_add").map_err(|e| {
        MLError::InitializationError {
            component: "CudaLinear".to_owned(),
            message: format!("Failed to load bias_add function: {e}"),
        }
    })
}

impl CudaLinear {
    /// Create a new linear layer with Xavier-uniform initialization.
    pub fn new(ctx: GpuContext, in_features: usize, out_features: usize) -> Result<Self, MLError> {
        let stream = &ctx.stream;

        // Xavier uniform: U(-sqrt(6 / (fan_in + fan_out)), sqrt(6 / (fan_in + fan_out)))
        let limit = (6.0_f32 / (in_features + out_features) as f32).sqrt();
        let weight_len = out_features * in_features;
        let mut weight_host = vec![0.0_f32; weight_len];
        let mut rng = fastrand::Rng::new();
        for w in &mut weight_host {
            *w = rng.f32() * 2.0 * limit - limit;
        }

        let mut weight = stream.alloc_zeros::<f32>(weight_len).map_err(|e| {
            MLError::InitializationError {
                component: "CudaLinear".to_owned(),
                message: format!("Failed to alloc weight: {e}"),
            }
        })?;
        stream.memcpy_htod(&weight_host, &mut weight).map_err(|e| {
            MLError::InitializationError {
                component: "CudaLinear".to_owned(),
                message: format!("Failed to upload weight: {e}"),
            }
        })?;

        let bias = stream.alloc_zeros::<f32>(out_features).map_err(|e| {
            MLError::InitializationError {
                component: "CudaLinear".to_owned(),
                message: format!("Failed to alloc bias: {e}"),
            }
        })?;

        let weight_grad = stream.alloc_zeros::<f32>(weight_len).map_err(|e| {
            MLError::InitializationError {
                component: "CudaLinear".to_owned(),
                message: format!("Failed to alloc weight_grad: {e}"),
            }
        })?;
        let bias_grad = stream.alloc_zeros::<f32>(out_features).map_err(|e| {
            MLError::InitializationError {
                component: "CudaLinear".to_owned(),
                message: format!("Failed to alloc bias_grad: {e}"),
            }
        })?;

        let bias_add_func = compile_bias_add(&ctx)?;

        Ok(Self {
            weight, bias, weight_grad, bias_grad,
            in_features, out_features, ctx, bias_add_func,
        })
    }

    /// Create from pre-existing weight data (for checkpoint loading).
    pub fn from_weights(
        ctx: GpuContext,
        in_features: usize,
        out_features: usize,
        weight_data: &[f32],
        bias_data: &[f32],
    ) -> Result<Self, MLError> {
        let stream = &ctx.stream;
        let weight_len = out_features * in_features;

        if weight_data.len() != weight_len {
            return Err(MLError::ConfigError(format!(
                "Weight data length {} != expected {}", weight_data.len(), weight_len
            )));
        }
        if bias_data.len() != out_features {
            return Err(MLError::ConfigError(format!(
                "Bias data length {} != expected {}", bias_data.len(), out_features
            )));
        }

        let mut weight = stream.alloc_zeros::<f32>(weight_len).map_err(|e| {
            MLError::InitializationError { component: "CudaLinear".to_owned(), message: format!("alloc weight: {e}") }
        })?;
        stream.memcpy_htod(weight_data, &mut weight).map_err(|e| {
            MLError::InitializationError { component: "CudaLinear".to_owned(), message: format!("upload weight: {e}") }
        })?;

        let mut bias = stream.alloc_zeros::<f32>(out_features).map_err(|e| {
            MLError::InitializationError { component: "CudaLinear".to_owned(), message: format!("alloc bias: {e}") }
        })?;
        stream.memcpy_htod(bias_data, &mut bias).map_err(|e| {
            MLError::InitializationError { component: "CudaLinear".to_owned(), message: format!("upload bias: {e}") }
        })?;

        let weight_grad = stream.alloc_zeros::<f32>(weight_len).map_err(|e| {
            MLError::InitializationError { component: "CudaLinear".to_owned(), message: format!("alloc weight_grad: {e}") }
        })?;
        let bias_grad = stream.alloc_zeros::<f32>(out_features).map_err(|e| {
            MLError::InitializationError { component: "CudaLinear".to_owned(), message: format!("alloc bias_grad: {e}") }
        })?;

        let bias_add_func = compile_bias_add(&ctx)?;

        Ok(Self {
            weight, bias, weight_grad, bias_grad,
            in_features, out_features, ctx, bias_add_func,
        })
    }

    /// Forward pass: `Y = X @ W^T + b`
    ///
    /// Input: `CudaSlice<f32>` of shape `[batch_size * in_features]` (row-major)
    /// Output: `CudaVec` of shape `[batch_size * out_features]` (row-major)
    pub fn forward(&self, input: &CudaSlice<f32>, batch_size: usize) -> Result<CudaVec, MLError> {
        let stream = &self.ctx.stream;
        let m = batch_size;
        let n = self.out_features;
        let k = self.in_features;

        let mut output = stream.alloc_zeros::<f32>(m * n).map_err(|e| {
            MLError::ModelError(format!("Failed to alloc linear output: {e}"))
        })?;

        // cuBLAS sgemm: Y = X @ W^T
        // Row-major → column-major convention:
        //   C[n,m] col = W[n,k] @ X^T[k,m]
        //   W stored as [n,k] row = [k,n] col → need CUBLAS_OP_T to get [n,k]
        //   X stored as [m,k] row = [k,m] col = X^T → CUBLAS_OP_N
        let w_ptr = raw_ptr(&self.weight, stream);
        let x_ptr = raw_ptr(input, stream);
        let y_ptr = raw_ptr_mut(&mut output, stream);

        // SAFETY: cuBLAS sgemm operates on valid device pointers from pre-allocated
        // weight [n,k], input [m,k], and output [m,n] CudaSlice buffers. All
        // dimensions are derived from the layer's construction parameters.
        unsafe {
            cudarc::cublas::result::sgemm(
                *self.ctx.cublas.handle(),
                cublasOperation_t::CUBLAS_OP_T,
                cublasOperation_t::CUBLAS_OP_N,
                n as i32, m as i32, k as i32,
                &1.0_f32 as *const f32,
                w_ptr as *const f32,
                k as i32,
                x_ptr as *const f32,
                k as i32,
                &0.0_f32 as *const f32,
                y_ptr as *mut f32,
                n as i32,
            ).map_err(|e| MLError::ModelError(format!("cuBLAS sgemm failed: {e}")))?;
        }

        // Add bias: Y[i,j] += bias[j]
        let total = m * n;
        let block_size = 256_u32;
        let grid_size = (total as u32).div_ceil(block_size);
        let n_i32 = n as i32;
        let total_i32 = total as i32;

        // SAFETY: bias_add kernel adds bias[j] to each output[i,j]. The output
        // buffer has `total` elements and bias has `n` elements, both pre-allocated.
        unsafe {
            self.ctx.stream
                .launch_builder(&self.bias_add_func)
                .arg(&mut output)
                .arg(&self.bias)
                .arg(&n_i32)
                .arg(&total_i32)
                .launch(LaunchConfig {
                    grid_dim: (grid_size, 1, 1),
                    block_dim: (block_size, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("bias_add kernel launch failed: {e}")))?;
        }

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

    /// Get total number of trainable parameters.
    pub const fn num_params(&self) -> usize {
        self.in_features * self.out_features + self.out_features
    }

    /// Zero the gradient buffers.
    pub fn zero_grad(&mut self) -> Result<(), MLError> {
        let stream = &self.ctx.stream;
        let w_zeros = vec![0.0_f32; self.in_features * self.out_features];
        stream.memcpy_htod(&w_zeros, &mut self.weight_grad).map_err(|e| {
            MLError::ModelError(format!("Failed to zero weight_grad: {e}"))
        })?;
        let b_zeros = vec![0.0_f32; self.out_features];
        stream.memcpy_htod(&b_zeros, &mut self.bias_grad).map_err(|e| {
            MLError::ModelError(format!("Failed to zero bias_grad: {e}"))
        })?;
        Ok(())
    }

    /// Download weights to host vectors (for checkpointing).
    #[cold]
    pub fn get_weights(&self) -> Result<(Vec<f32>, Vec<f32>), MLError> {
        let stream = &self.ctx.stream;
        let weight_len = self.in_features * self.out_features;
        let mut weight_host = vec![0.0_f32; weight_len];
        stream.memcpy_dtoh(&self.weight, &mut weight_host).map_err(|e| { // gpu-exit: checkpoint weight export
            MLError::ModelError(format!("Failed to download weight: {e}"))
        })?;
        let mut bias_host = vec![0.0_f32; self.out_features];
        stream.memcpy_dtoh(&self.bias, &mut bias_host).map_err(|e| { // gpu-exit: checkpoint weight export
            MLError::ModelError(format!("Failed to download bias: {e}"))
        })?;
        Ok((weight_host, bias_host))
    }

    /// Upload weights from host vectors (for checkpoint loading).
    #[cold]
    pub fn set_weights(&mut self, weights: &[f32], bias: &[f32]) -> Result<(), MLError> {
        let expected_w = self.in_features * self.out_features;
        if weights.len() != expected_w {
            return Err(MLError::DimensionMismatch {
                expected: expected_w,
                actual: weights.len(),
            });
        }
        if bias.len() != self.out_features {
            return Err(MLError::DimensionMismatch {
                expected: self.out_features,
                actual: bias.len(),
            });
        }
        let stream = &self.ctx.stream;
        stream.memcpy_htod(weights, &mut self.weight).map_err(|e| { // gpu-entry: checkpoint weight import
            MLError::ModelError(format!("Failed to upload weight: {e}"))
        })?;
        stream.memcpy_htod(bias, &mut self.bias).map_err(|e| { // gpu-entry: checkpoint weight import
            MLError::ModelError(format!("Failed to upload bias: {e}"))
        })?;
        Ok(())
    }
}
