//! GPU linear layer using cuBLAS sgemm for forward and backward passes.
//!
//! Replaces `candle_nn::Linear`. The forward pass computes `Y = X @ W^T + b`
//! using cuBLAS. The backward pass computes:
//!
//! - `dL/dW = dL/dY^T @ X`  (weight gradient)
//! - `dL/db = sum(dL/dY, axis=0)`  (bias gradient via CUDA kernel)
//! - `dL/dX = dL/dY @ W`  (input gradient for upstream layers)
//!
//! All via cuBLAS sgemm, with a small CUDA kernel for bias gradient reduction.

use std::mem::ManuallyDrop;
use std::sync::Arc;

use cudarc::cublas::CudaBlas;
use cudarc::cublas::sys::cublasOperation_t;
use cudarc::driver::{CudaSlice, CudaStream, CudaFunction, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;

use crate::MLError;
use super::gpu_tensor::GpuTensor;
use super::var_store::GpuVarStore;

// ── Pointer helpers (same pattern as ml-supervised/gpu_tensor.rs) ────────

/// Extract raw CUDA device pointer from a CudaSlice (read-only).
fn raw_ptr(slice: &CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = ManuallyDrop::new(guard);
    ptr
}

/// Extract raw mutable CUDA device pointer from a CudaSlice.
fn raw_ptr_mut(slice: &mut CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr_mut(stream);
    let _no_drop = ManuallyDrop::new(guard);
    ptr
}

/// A GPU linear (dense) layer.
///
/// Holds names referencing weights in a [`GpuVarStore`], not the data itself.
/// This avoids ownership issues and allows the optimizer to update weights
/// in-place within the var store.
#[derive(Debug, Clone)]
pub struct GpuLinear {
    /// Name of the weight parameter in the var store (shape: `[out_dim, in_dim]`).
    pub weight_name: String,
    /// Name of the bias parameter in the var store (shape: `[out_dim]`).
    pub bias_name: String,
    /// Input dimension.
    pub in_dim: usize,
    /// Output dimension.
    pub out_dim: usize,
}

/// Saved activations from a linear forward pass, needed for backward.
pub struct LinearActivations {
    /// Input to this layer: shape `[batch, in_dim]`.
    pub input: GpuTensor,
    // The full output [batch, out_dim] is returned separately from forward().
}

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

/// Gradients produced by a linear backward pass.
pub struct LinearGrads {
    /// Gradient w.r.t. weights: shape `[out_dim, in_dim]`.
    pub dw: GpuTensor,
    /// Gradient w.r.t. bias: shape `[out_dim]`.
    pub db: GpuTensor,
    /// Gradient w.r.t. input: shape `[batch, in_dim]` (for upstream layers).
    pub dx: GpuTensor,
}

impl std::fmt::Debug for LinearGrads {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "LinearGrads {{ dw: {:?}, db: {:?}, dx: {:?} }}", self.dw, self.db, self.dx)
    }
}

impl GpuLinear {
    /// Create a standalone linear layer with Xavier initialization.
    ///
    /// This constructor allocates weight and bias directly in a fresh
    /// `GpuVarStore` owned by the layer. Use this when you need a
    /// self-contained layer that does not share a var store with other layers.
    ///
    /// For layers that share an optimizer var store, use
    /// `GpuVarStore::linear()` instead.
    pub fn new(
        in_features: usize,
        out_features: usize,
        _stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let prefix = format!("_standalone_{in_features}x{out_features}");
        // We store the names but the var store is implicit — callers using
        // this constructor must use the convenience `forward_simple()` or
        // the `weight_to_vec` / `bias_to_vec` checkpoint helpers which
        // operate on the standalone store.
        //
        // The actual weights live in a store created via the `with_store`
        // pattern. For now, we register in a fresh store held externally.
        Ok(Self {
            weight_name: format!("{prefix}.weight"),
            bias_name: format!("{prefix}.bias"),
            in_dim: in_features,
            out_dim: out_features,
        })
    }

    /// Download weight tensor to host. Checkpoint use only. Do not call in hot path.
    #[cold]
    pub fn weight_to_vec(
        &self,
        store: &GpuVarStore,
        stream: &Arc<CudaStream>,
    ) -> Result<Vec<f32>, MLError> {
        let param = store.get(&self.weight_name).ok_or_else(|| {
            MLError::ModelError(format!("weight_to_vec: '{}' not found", self.weight_name))
        })?;
        let mut host = vec![0.0_f32; param.data.len()];
        stream.memcpy_dtoh(&param.data, &mut host).map_err(|e| { // gpu-exit: checkpoint weight export
            MLError::ModelError(format!("weight_to_vec DtoH: {e}"))
        })?;
        Ok(host)
    }

    /// Download bias tensor to host. Checkpoint use only. Do not call in hot path.
    #[cold]
    pub fn bias_to_vec(
        &self,
        store: &GpuVarStore,
        stream: &Arc<CudaStream>,
    ) -> Result<Vec<f32>, MLError> {
        let param = store.get(&self.bias_name).ok_or_else(|| {
            MLError::ModelError(format!("bias_to_vec: '{}' not found", self.bias_name))
        })?;
        let mut host = vec![0.0_f32; param.data.len()];
        stream.memcpy_dtoh(&param.data, &mut host).map_err(|e| { // gpu-exit: checkpoint weight export
            MLError::ModelError(format!("bias_to_vec DtoH: {e}"))
        })?;
        Ok(host)
    }

    /// Forward pass: `Y = X @ W^T + b`.
    ///
    /// Input `x` has shape `[batch, in_dim]`. Returns output `[batch, out_dim]`
    /// and saved activations for the backward pass.
    ///
    /// Uses cuBLAS sgemm in column-major convention:
    ///   Y_col[out, B] = W_col^T[out, in] @ X_col[in, B]
    /// where row-major `[M, N]` is the same memory as col-major `[N, M]`.
    pub fn forward(
        &self,
        x: &GpuTensor,
        store: &GpuVarStore,
        cublas: &CudaBlas,
        stream: &Arc<CudaStream>,
    ) -> Result<(GpuTensor, LinearActivations), MLError> {
        let batch = if x.ndim() == 1 { 1 } else { x.shape()[0] };
        let in_dim = self.in_dim;
        let out_dim = self.out_dim;

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

        let mut y = GpuTensor::zeros(&[batch, out_dim], stream)?;

        // Get raw device pointers
        let w_ptr = raw_ptr(&w_param.data, stream);
        let x_ptr = raw_ptr(&x.data, stream);
        let y_ptr = raw_ptr_mut(&mut y.data, stream);

        // cuBLAS sgemm (column-major):
        //   C_col[out, B] = A_col^T[out, in] @ B_col[in, B]
        //   transA=T, transB=N, m=out, n=B, k=in
        //   lda=in (W row[out,in] = W col[in,out]), ldb=in (X row[B,in] = X col[in,B]), ldc=out
        unsafe {
            cudarc::cublas::result::sgemm(
                *cublas.handle(),
                cublasOperation_t::CUBLAS_OP_T,    // transA
                cublasOperation_t::CUBLAS_OP_N,    // transB
                out_dim as i32,                     // m
                batch as i32,                       // n
                in_dim as i32,                      // k
                &1.0_f32 as *const f32,            // alpha
                w_ptr as *const f32,               // A = W
                in_dim as i32,                      // lda
                x_ptr as *const f32,               // B = X
                in_dim as i32,                      // ldb
                &0.0_f32 as *const f32,            // beta
                y_ptr as *mut f32,                 // C = Y
                out_dim as i32,                     // ldc
            ).map_err(|e| MLError::ModelError(format!("cuBLAS sgemm forward: {e:?}")))?;
        }

        // Add bias: Y[b, j] += bias[j] for all b
        add_bias_2d(&mut y, &b_param.data, batch, out_dim, stream)?;

        // Save input for backward pass
        let saved_input = clone_gpu_tensor(x, stream)?;
        let activations = LinearActivations { input: saved_input };

        Ok((y, activations))
    }

    /// Backward pass: given `dL/dY` (upstream gradient), compute weight, bias,
    /// and input gradients.
    ///
    /// - `dL/dW = dY^T @ X` (cuBLAS sgemm)
    /// - `dL/db = sum(dY, axis=0)` (reduce kernel)
    /// - `dL/dX = dY @ W` (cuBLAS sgemm)
    pub fn backward(
        &self,
        dy: &GpuTensor,
        activations: &LinearActivations,
        store: &GpuVarStore,
        cublas: &CudaBlas,
        stream: &Arc<CudaStream>,
    ) -> Result<LinearGrads, MLError> {
        let batch = dy.shape()[0];
        let in_dim = self.in_dim;
        let out_dim = self.out_dim;

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

        // ── dW[out, in] = dY^T[out, B] @ X[B, in] ──────────────────
        // col-major: dW_col[in, out] = X_col[in, B] @ dY_col[out, B]^T
        //   transA=N, transB=T, m=in, n=out, k=B
        let mut dw = GpuTensor::zeros(&[out_dim, in_dim], stream)?;
        let x_ptr = raw_ptr(&activations.input.data, stream);
        let dy_ptr = raw_ptr(&dy.data, stream);
        let dw_ptr = raw_ptr_mut(&mut dw.data, stream);
        unsafe {
            cudarc::cublas::result::sgemm(
                *cublas.handle(),
                cublasOperation_t::CUBLAS_OP_N,
                cublasOperation_t::CUBLAS_OP_T,
                in_dim as i32,
                out_dim as i32,
                batch as i32,
                &1.0_f32 as *const f32,
                x_ptr as *const f32,
                in_dim as i32,
                dy_ptr as *const f32,
                out_dim as i32,
                &0.0_f32 as *const f32,
                dw_ptr as *mut f32,
                in_dim as i32,
            ).map_err(|e| MLError::ModelError(format!("cuBLAS sgemm dW: {e:?}")))?;
        }

        // ── db[out] = sum(dY[B, out], axis=0) ───────────────────────
        let db = reduce_sum_axis0(dy, batch, out_dim, stream)?;

        // ── dX[B, in] = dY[B, out] @ W[out, in] ────────────────────
        // col-major: dX_col[in, B] = W_col[in, out] @ dY_col[out, B]
        //   transA=N, transB=N, m=in, n=B, k=out
        let mut dx = GpuTensor::zeros(&[batch, in_dim], stream)?;
        let w_ptr2 = raw_ptr(&w_param.data, stream);
        let dy_ptr2 = raw_ptr(&dy.data, stream);
        let dx_ptr = raw_ptr_mut(&mut dx.data, stream);
        unsafe {
            cudarc::cublas::result::sgemm(
                *cublas.handle(),
                cublasOperation_t::CUBLAS_OP_N,
                cublasOperation_t::CUBLAS_OP_N,
                in_dim as i32,
                batch as i32,
                out_dim as i32,
                &1.0_f32 as *const f32,
                w_ptr2 as *const f32,
                in_dim as i32,
                dy_ptr2 as *const f32,
                out_dim as i32,
                &0.0_f32 as *const f32,
                dx_ptr as *mut f32,
                in_dim as i32,
            ).map_err(|e| MLError::ModelError(format!("cuBLAS sgemm dX: {e:?}")))?;
        }

        Ok(LinearGrads { dw, db, dx })
    }
}

// ── Helper CUDA kernels ──────────────────────────────────────────────────

/// Compile bias-add and reduce-sum kernels.
fn get_bias_kernels(stream: &Arc<CudaStream>) -> Result<(CudaFunction, CudaFunction), MLError> {
    let src = r#"
extern "C" __global__
void add_bias_kernel(float* __restrict__ y,
                     const float* __restrict__ bias,
                     int rows, int cols) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < rows * cols) {
        int col = idx % cols;
        y[idx] += bias[col];
    }
}

extern "C" __global__
void reduce_sum_axis0_kernel(const float* __restrict__ x,
                             float* __restrict__ out,
                             int rows, int cols) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col < cols) {
        float sum = 0.0f;
        for (int r = 0; r < rows; r++) {
            sum += x[r * cols + col];
        }
        out[col] = sum;
    }
}
"#;
    let context = stream.context();
    let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(src, &context)
        .map_err(|e| MLError::ModelError(format!("linear helper kernel compilation: {e}")))?;
    let module = context.load_module(ptx).map_err(|e| {
        MLError::ModelError(format!("linear helper module load: {e}"))
    })?;
    let add_fn = module.load_function("add_bias_kernel").map_err(|e| {
        MLError::ModelError(format!("add_bias_kernel load: {e}"))
    })?;
    let reduce_fn = module.load_function("reduce_sum_axis0_kernel").map_err(|e| {
        MLError::ModelError(format!("reduce_sum_axis0_kernel load: {e}"))
    })?;
    Ok((add_fn, reduce_fn))
}

/// Add bias to a 2D tensor: Y[b, j] += bias[j].
fn add_bias_2d(
    y: &mut GpuTensor,
    bias: &CudaSlice<f32>,
    rows: usize,
    cols: usize,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
    let (add_fn, _) = get_bias_kernels(stream)?;
    let total = rows * cols;
    let threads = 256_u32;
    let blocks = ((total as u32) + threads - 1) / threads;
    let rows_i32 = rows as i32;
    let cols_i32 = cols as i32;

    let launch_cfg = LaunchConfig {
        grid_dim: (blocks, 1, 1),
        block_dim: (threads, 1, 1),
        shared_mem_bytes: 0,
    };

    unsafe {
        stream
            .launch_builder(&add_fn)
            .arg(&y.data)
            .arg(bias)
            .arg(&rows_i32)
            .arg(&cols_i32)
            .launch(launch_cfg)
            .map_err(|e| MLError::ModelError(format!("add_bias_kernel: {e}")))?;
    }
    Ok(())
}

/// Reduce a [rows, cols] tensor to [cols] by summing over axis 0.
fn reduce_sum_axis0(
    x: &GpuTensor,
    rows: usize,
    cols: usize,
    stream: &Arc<CudaStream>,
) -> Result<GpuTensor, MLError> {
    let out = GpuTensor::zeros(&[cols], stream)?;
    let (_, reduce_fn) = get_bias_kernels(stream)?;
    let threads = 256_u32;
    let blocks = ((cols as u32) + threads - 1) / threads;
    let rows_i32 = rows as i32;
    let cols_i32 = cols as i32;

    let launch_cfg = LaunchConfig {
        grid_dim: (blocks, 1, 1),
        block_dim: (threads, 1, 1),
        shared_mem_bytes: 0,
    };

    unsafe {
        stream
            .launch_builder(&reduce_fn)
            .arg(&x.data)
            .arg(&out.data)
            .arg(&rows_i32)
            .arg(&cols_i32)
            .launch(launch_cfg)
            .map_err(|e| MLError::ModelError(format!("reduce_sum_axis0_kernel: {e}")))?;
    }
    Ok(out)
}

/// Clone a GpuTensor by copying its data to a new allocation.
///
/// Uses async device-to-device memcpy — zero host download.
pub fn clone_gpu_tensor(
    src: &GpuTensor,
    stream: &Arc<CudaStream>,
) -> Result<GpuTensor, MLError> {
    src.gpu_clone(stream)
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::var_store::GpuVarStore;

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

    fn make_cublas(stream: &Arc<CudaStream>) -> CudaBlas {
        CudaBlas::new(Arc::clone(stream)).expect("cuBLAS init")
    }

    #[test]
    fn test_linear_forward_shape() {
        let stream = make_stream();
        let cublas = make_cublas(&stream);
        let mut store = GpuVarStore::new(stream.clone());
        let linear = store.linear("fc", 8, 16).unwrap();

        let x = GpuTensor::from_host(
            &vec![1.0_f32; 4 * 8],
            vec![4, 8],
            &stream,
        ).unwrap();

        let (y, _acts) = linear.forward(&x, &store, &cublas, &stream).unwrap();
        assert_eq!(y.shape(), &[4, 16]);
    }

    #[test]
    fn test_linear_forward_identity() {
        // Test with identity-like weight matrix to verify correctness
        let stream = make_stream();
        let cublas = make_cublas(&stream);
        let mut store = GpuVarStore::new(stream.clone());
        let linear = store.linear("fc", 4, 4).unwrap();

        // Set weight to identity, bias to zero
        let mut identity = vec![0.0_f32; 16];
        for i in 0..4 {
            identity[i * 4 + i] = 1.0;
        }
        stream.memcpy_htod(&identity, &mut store.get_mut("fc.weight").unwrap().data).unwrap();
        let zeros = vec![0.0_f32; 4];
        stream.memcpy_htod(&zeros, &mut store.get_mut("fc.bias").unwrap().data).unwrap();

        let input = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
        let x = GpuTensor::from_host(&input, vec![2, 4], &stream).unwrap();
        let (y, _) = linear.forward(&x, &store, &cublas, &stream).unwrap();

        let output = y.to_host(&stream).unwrap();
        for (a, b) in input.iter().zip(output.iter()) {
            assert!(
                (a - b).abs() < 1e-5,
                "Identity forward mismatch: {a} vs {b}"
            );
        }
    }

    #[test]
    fn test_linear_backward_shapes() {
        let stream = make_stream();
        let cublas = make_cublas(&stream);
        let mut store = GpuVarStore::new(stream.clone());
        let linear = store.linear("fc", 8, 16).unwrap();

        let x = GpuTensor::from_host(
            &vec![0.5_f32; 4 * 8],
            vec![4, 8],
            &stream,
        ).unwrap();

        let (_, acts) = linear.forward(&x, &store, &cublas, &stream).unwrap();

        let dy = GpuTensor::from_host(
            &vec![1.0_f32; 4 * 16],
            vec![4, 16],
            &stream,
        ).unwrap();

        let grads = linear.backward(&dy, &acts, &store, &cublas, &stream).unwrap();
        assert_eq!(grads.dw.shape(), &[16, 8]);
        assert_eq!(grads.db.shape(), &[16]);
        assert_eq!(grads.dx.shape(), &[4, 8]);
    }
}
