//! GPU-resident LSTM layer using cuBLAS for matrix multiplications.
//!
//! Implements a standard LSTM cell with fused gate computation.
//! All gates computed in a single cuBLAS sgemm + element-wise 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 LSTM gate kernel -- compiled once per process via `OnceLock`.
static LSTM_GATE_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();
/// Cached PTX for LSTM `bias_add` kernel -- compiled once per process via `OnceLock`.
static LSTM_BIAS_ADD_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();

/// CUDA LSTM gate fusion kernel.
const LSTM_GATE_KERNEL: &str = r#"
extern "C" __global__ void lstm_gates(
    float* __restrict__ h_new,
    float* __restrict__ c_new,
    const float* __restrict__ gates,
    const float* __restrict__ c_old,
    int hidden_dim,
    int batch_size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * hidden_dim;
    if (idx >= total) return;

    int b = idx / hidden_dim;
    int h = idx % hidden_dim;
    int gates_offset = b * 4 * hidden_dim;

    float i_gate = 1.0f / (1.0f + expf(-gates[gates_offset + h]));
    float f_gate = 1.0f / (1.0f + expf(-gates[gates_offset + hidden_dim + h]));
    float g_gate = tanhf(gates[gates_offset + 2 * hidden_dim + h]);
    float o_gate = 1.0f / (1.0f + expf(-gates[gates_offset + 3 * hidden_dim + h]));

    float c_val = f_gate * c_old[idx] + i_gate * g_gate;
    float h_val = o_gate * tanhf(c_val);

    c_new[idx] = c_val;
    h_new[idx] = h_val;
}
"#;

const BIAS_ADD_KERNEL: &str = r#"
extern "C" __global__ void bias_add_lstm(
    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];
    }
}
"#;

/// GPU-resident LSTM layer.
pub struct CudaLSTM {
    /// Input-to-hidden weights: `[4*hidden_dim, input_dim]` (row-major)
    pub w_ih: CudaSlice<f32>,
    /// Hidden-to-hidden weights: `[4*hidden_dim, hidden_dim]` (row-major)
    pub w_hh: CudaSlice<f32>,
    /// Input-to-hidden bias: `[4*hidden_dim]`
    pub b_ih: CudaSlice<f32>,
    /// Hidden-to-hidden bias: `[4*hidden_dim]`
    pub b_hh: CudaSlice<f32>,
    /// Gradient buffers
    pub w_ih_grad: CudaSlice<f32>,
    pub w_hh_grad: CudaSlice<f32>,
    pub b_ih_grad: CudaSlice<f32>,
    pub b_hh_grad: CudaSlice<f32>,
    /// Dimensions
    pub input_dim: usize,
    pub hidden_dim: usize,
    /// CUDA context
    ctx: GpuContext,
    /// Gate kernel
    gate_func: cudarc::driver::CudaFunction,
    /// Bias-add kernel
    bias_add_func: cudarc::driver::CudaFunction,
}

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

impl CudaLSTM {
    /// Create a new LSTM layer with Xavier initialization.
    pub fn new(ctx: GpuContext, input_dim: usize, hidden_dim: usize) -> Result<Self, MLError> {
        let stream = &ctx.stream;
        let gate_dim = 4 * hidden_dim;

        let limit_ih = (6.0_f32 / (input_dim + hidden_dim) as f32).sqrt();
        let limit_hh = (6.0_f32 / (hidden_dim + hidden_dim) as f32).sqrt();

        let w_ih_len = gate_dim * input_dim;
        let w_hh_len = gate_dim * hidden_dim;

        let mut rng = fastrand::Rng::new();

        let w_ih_host: Vec<f32> = (0..w_ih_len).map(|_| rng.f32() * 2.0 * limit_ih - limit_ih).collect();
        let w_hh_host: Vec<f32> = (0..w_hh_len).map(|_| rng.f32() * 2.0 * limit_hh - limit_hh).collect();

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

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

        let b_ih = stream.alloc_zeros::<f32>(gate_dim).map_err(|e| {
            MLError::InitializationError { component: "CudaLSTM".to_owned(), message: format!("alloc b_ih: {e}") }
        })?;
        let b_hh = stream.alloc_zeros::<f32>(gate_dim).map_err(|e| {
            MLError::InitializationError { component: "CudaLSTM".to_owned(), message: format!("alloc b_hh: {e}") }
        })?;

        let w_ih_grad = stream.alloc_zeros::<f32>(w_ih_len).map_err(|e| {
            MLError::InitializationError { component: "CudaLSTM".to_owned(), message: format!("alloc w_ih_grad: {e}") }
        })?;
        let w_hh_grad = stream.alloc_zeros::<f32>(w_hh_len).map_err(|e| {
            MLError::InitializationError { component: "CudaLSTM".to_owned(), message: format!("alloc w_hh_grad: {e}") }
        })?;
        let b_ih_grad = stream.alloc_zeros::<f32>(gate_dim).map_err(|e| {
            MLError::InitializationError { component: "CudaLSTM".to_owned(), message: format!("alloc b_ih_grad: {e}") }
        })?;
        let b_hh_grad = stream.alloc_zeros::<f32>(gate_dim).map_err(|e| {
            MLError::InitializationError { component: "CudaLSTM".to_owned(), message: format!("alloc b_hh_grad: {e}") }
        })?;

        let context = stream.context();

        let gate_ptx_result = LSTM_GATE_PTX.get_or_init(|| {
            ml_core::cuda_compile::compile_ptx_for_device(LSTM_GATE_KERNEL, context)
                .map_err(|e| format!("compile gate kernel: {e}"))
        });
        let gate_ptx = gate_ptx_result.as_ref().map_err(|e| {
            MLError::InitializationError { component: "CudaLSTM".to_owned(), message: e.clone() }
        })?;
        let gate_module = context.load_module(gate_ptx.clone()).map_err(|e| {
            MLError::InitializationError { component: "CudaLSTM".to_owned(), message: format!("load gate module: {e}") }
        })?;
        let gate_func = gate_module.load_function("lstm_gates").map_err(|e| {
            MLError::InitializationError { component: "CudaLSTM".to_owned(), message: format!("load lstm_gates: {e}") }
        })?;

        let bias_ptx_result = LSTM_BIAS_ADD_PTX.get_or_init(|| {
            ml_core::cuda_compile::compile_ptx_for_device(BIAS_ADD_KERNEL, context)
                .map_err(|e| format!("compile bias_add: {e}"))
        });
        let bias_ptx = bias_ptx_result.as_ref().map_err(|e| {
            MLError::InitializationError { component: "CudaLSTM".to_owned(), message: e.clone() }
        })?;
        let bias_module = context.load_module(bias_ptx.clone()).map_err(|e| {
            MLError::InitializationError { component: "CudaLSTM".to_owned(), message: format!("load bias_add: {e}") }
        })?;
        let bias_add_func = bias_module.load_function("bias_add_lstm").map_err(|e| {
            MLError::InitializationError { component: "CudaLSTM".to_owned(), message: format!("load bias_add_lstm: {e}") }
        })?;

        Ok(Self {
            w_ih, w_hh, b_ih, b_hh,
            w_ih_grad, w_hh_grad, b_ih_grad, b_hh_grad,
            input_dim, hidden_dim,
            ctx, gate_func, bias_add_func,
        })
    }

    /// LSTM forward step.
    ///
    /// Input: x `[batch, input_dim]`, h `[batch, hidden_dim]`, c `[batch, hidden_dim]`.
    /// Output: (`h_new` `[batch, hidden_dim]`, `c_new` `[batch, hidden_dim]`)
    pub fn forward(
        &self,
        x: &CudaSlice<f32>,
        h: &CudaSlice<f32>,
        c: &CudaSlice<f32>,
        batch_size: usize,
    ) -> Result<(CudaVec, CudaVec), MLError> {
        let stream = &self.ctx.stream;
        let gate_dim = 4 * self.hidden_dim;

        let mut gates = stream.alloc_zeros::<f32>(batch_size * gate_dim).map_err(|e| {
            MLError::ModelError(format!("alloc gates: {e}"))
        })?;

        // gates = W_ih @ x^T (via cuBLAS)
        let w_ih_ptr = raw_ptr(&self.w_ih, stream);
        let x_ptr = raw_ptr(x, stream);
        let gates_ptr = raw_ptr_mut(&mut gates, stream);

        // SAFETY: cuBLAS sgemm computes gates = W_ih @ x^T. All device pointers
        // are valid and derived from pre-allocated CudaSlice buffers with dimensions
        // matching the W_ih [gate_dim, input_dim] and x [batch, input_dim] shapes.
        unsafe {
            cudarc::cublas::result::sgemm(
                *self.ctx.cublas.handle(),
                cublasOperation_t::CUBLAS_OP_T,
                cublasOperation_t::CUBLAS_OP_N,
                gate_dim as i32,
                batch_size as i32,
                self.input_dim as i32,
                &1.0_f32,
                w_ih_ptr as *const f32,
                self.input_dim as i32,
                x_ptr as *const f32,
                self.input_dim as i32,
                &0.0_f32,
                gates_ptr as *mut f32,
                gate_dim as i32,
            ).map_err(|e| MLError::ModelError(format!("sgemm W_ih @ x: {e}")))?;
        }

        // gates += b_ih (broadcast)
        let total_gates = batch_size * gate_dim;
        let block_size = 256_u32;
        let grid_size = (total_gates as u32).div_ceil(block_size);
        let gate_dim_i32 = gate_dim as i32;
        let total_i32 = total_gates as i32;
        // SAFETY: bias_add kernel broadcasts b_ih [gate_dim] across
        // total_gates elements in the gates buffer. Both are pre-allocated.
        unsafe {
            stream.launch_builder(&self.bias_add_func)
                .arg(&mut gates)
                .arg(&self.b_ih)
                .arg(&gate_dim_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 b_ih: {e}")))?;
        }

        // gates += W_hh @ h^T + b_hh
        let w_hh_ptr = raw_ptr(&self.w_hh, stream);
        let h_ptr = raw_ptr(h, stream);
        let gates_ptr2 = raw_ptr_mut(&mut gates, stream);

        // SAFETY: cuBLAS sgemm accumulates gates += W_hh @ h^T (beta=1). All device
        // pointers are valid from pre-allocated W_hh [gate_dim, hidden_dim] and
        // h [batch, hidden_dim] buffers.
        unsafe {
            cudarc::cublas::result::sgemm(
                *self.ctx.cublas.handle(),
                cublasOperation_t::CUBLAS_OP_T,
                cublasOperation_t::CUBLAS_OP_N,
                gate_dim as i32,
                batch_size as i32,
                self.hidden_dim as i32,
                &1.0_f32,
                w_hh_ptr as *const f32,
                self.hidden_dim as i32,
                h_ptr as *const f32,
                self.hidden_dim as i32,
                &1.0_f32,  // beta=1 to accumulate
                gates_ptr2 as *mut f32,
                gate_dim as i32,
            ).map_err(|e| MLError::ModelError(format!("sgemm W_hh @ h: {e}")))?;
        }

        // gates += b_hh
        // SAFETY: bias_add kernel broadcasts b_hh [gate_dim] across total_gates
        // elements in the gates buffer. Both are pre-allocated with matching sizes.
        unsafe {
            stream.launch_builder(&self.bias_add_func)
                .arg(&mut gates)
                .arg(&self.b_hh)
                .arg(&gate_dim_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 b_hh: {e}")))?;
        }

        // Apply gate nonlinearities and compute h_new, c_new
        let hc_total = batch_size * self.hidden_dim;
        let mut h_new = stream.alloc_zeros::<f32>(hc_total).map_err(|e| {
            MLError::ModelError(format!("alloc h_new: {e}"))
        })?;
        let mut c_new = stream.alloc_zeros::<f32>(hc_total).map_err(|e| {
            MLError::ModelError(format!("alloc c_new: {e}"))
        })?;

        let hc_grid = (hc_total as u32).div_ceil(block_size);
        let hidden_i32 = self.hidden_dim as i32;
        let batch_i32 = batch_size as i32;

        // SAFETY: lstm_gate_kernel applies sigmoid/tanh nonlinearities to the
        // combined gate values and computes h_new/c_new. All buffers (gates,
        // c, h_new, c_new) are pre-allocated with correct dimensions.
        unsafe {
            stream.launch_builder(&self.gate_func)
                .arg(&mut h_new)
                .arg(&mut c_new)
                .arg(&gates)
                .arg(c)
                .arg(&hidden_i32)
                .arg(&batch_i32)
                .launch(LaunchConfig {
                    grid_dim: (hc_grid, 1, 1),
                    block_dim: (block_size, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("lstm_gates launch: {e}")))?;
        }

        Ok((
            CudaVec::new(h_new, hc_total),
            CudaVec::new(c_new, hc_total),
        ))
    }

    /// Get total number of trainable parameters.
    pub const fn num_params(&self) -> usize {
        let gate_dim = 4 * self.hidden_dim;
        gate_dim * self.input_dim + gate_dim * self.hidden_dim + 2 * gate_dim
    }
}
