//! sLSTM cell -- scalar LSTM with exponential gating.
//!
//! Unlike standard LSTM which uses sigmoid gates, sLSTM uses exponential gates
//! for the input and forget gates: `exp(W_i * x)` instead of `sigmoid(W_i * x)`.
//! This prevents gradient vanishing over long sequences. A normalizer state `n_t`
//! stabilizes the exponential values.

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::MLError;

use crate::gpu_tensor::{
    gpu_add, gpu_clamp, gpu_exp, gpu_mul, gpu_sigmoid, gpu_tanh, GpuLinear, GpuTensor,
};

/// sLSTM cell with exponential gating and scalar memory.
#[derive(Debug)]
pub struct SLSTMCell {
    /// Combined gate projection: x -> (input, forget, output, `cell_input`) * 4
    gate_proj: GpuLinear,
    /// Hidden-to-hidden recurrence
    hidden_proj: GpuLinear,
    hidden_dim: usize,
}

impl SLSTMCell {
    /// Create a new sLSTM cell.
    ///
    /// * `input_dim` - Dimension of input features
    /// * `hidden_dim` - Dimension of hidden state and cell state
    /// * `stream` - CUDA stream
    pub fn new(input_dim: usize, hidden_dim: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        // Projects input to 4 gates: i, f, o, z (cell input)
        let gate_proj = GpuLinear::new(input_dim, 4 * hidden_dim, stream)?;
        let hidden_proj = GpuLinear::new(hidden_dim, 4 * hidden_dim, stream)?;

        Ok(Self {
            gate_proj,
            hidden_proj,
            hidden_dim,
        })
    }

    /// Forward pass through sLSTM cell.
    ///
    /// * `x` - Input tensor of shape `(batch, input_dim)`
    /// * `state` - Optional `(h, c)` where h and c are `(batch, hidden_dim)`.
    ///   If None, zero-initializes both.
    ///
    /// Returns `(h_new, c_new)` both of shape `(batch, hidden_dim)`.
    pub fn forward(
        &self,
        x: &GpuTensor,
        state: Option<(&GpuTensor, &GpuTensor)>,
    ) -> Result<(GpuTensor, GpuTensor), MLError> {
        let batch = x.dim(0)?;

        let (h_prev, c_prev) = if let Some((h, c)) = state {
            (h.clone(), c.clone())
        } else {
            let h_zeros = GpuTensor::zeros(&[batch, self.hidden_dim], &x.stream)?;
            let c_zeros = GpuTensor::zeros(&[batch, self.hidden_dim], &x.stream)?;
            (h_zeros, c_zeros)
        };

        // Combined projection: gates = W_x * x + W_h * h_prev
        let x_proj = self.gate_proj.forward(x)?;
        let h_proj = self.hidden_proj.forward(&h_prev)?;
        let gates = gpu_add(&x_proj, &h_proj)?;

        // Split into 4 gates along last dimension
        let hd = self.hidden_dim;
        let i_gate = crate::gpu_tensor::gpu_narrow_2d(&gates, 1, 0, hd)?;
        let f_gate = crate::gpu_tensor::gpu_narrow_2d(&gates, 1, hd, hd)?;
        let o_gate = crate::gpu_tensor::gpu_narrow_2d(&gates, 1, 2 * hd, hd)?;
        let z_gate = crate::gpu_tensor::gpu_narrow_2d(&gates, 1, 3 * hd, hd)?;

        // Exponential gating (clamped to avoid overflow)
        let i_clamped = gpu_clamp(&i_gate, -15.0, 15.0)?;
        let i_exp = gpu_exp(&i_clamped)?;
        let f_clamped = gpu_clamp(&f_gate, -15.0, 15.0)?;
        let f_exp = gpu_exp(&f_clamped)?;

        // Output gate: standard sigmoid
        let o_sig = gpu_sigmoid(&o_gate)?;

        // Cell input: tanh
        let z_tanh = gpu_tanh(&z_gate)?;

        // Cell update: c_t = f_t * c_{t-1} + i_t * z_t
        let fc = gpu_mul(&f_exp, &c_prev)?;
        let iz = gpu_mul(&i_exp, &z_tanh)?;
        let c_new = gpu_add(&fc, &iz)?;

        // h_t = o_t * tanh(c_t)
        let c_tanh = gpu_tanh(&c_new)?;
        let h_new = gpu_mul(&o_sig, &c_tanh)?;

        Ok((h_new, c_new))
    }

    /// Return hidden dimension.
    pub const fn hidden_dim(&self) -> usize {
        self.hidden_dim
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cudarc::driver::CudaContext;

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA context required");
        ctx.new_stream().expect("Failed to create CUDA stream")
    }

    #[test]
    fn test_slstm_output_shape() {
        let stream = test_stream();
        let cell = SLSTMCell::new(32, 16, &stream).unwrap();

        let x = GpuTensor::randn(&[4, 32], 1.0, &stream).unwrap();
        let (h, c) = cell.forward(&x, None).unwrap();
        assert_eq!(h.shape, vec![4, 16]);
        assert_eq!(c.shape, vec![4, 16]);
    }

    #[test]
    fn test_slstm_sequential() {
        let stream = test_stream();
        let cell = SLSTMCell::new(32, 16, &stream).unwrap();

        let x1 = GpuTensor::randn(&[4, 32], 1.0, &stream).unwrap();
        let (h1, c1) = cell.forward(&x1, None).unwrap();

        let x2 = GpuTensor::randn(&[4, 32], 1.0, &stream).unwrap();
        let (h2, _c2) = cell.forward(&x2, Some((&h1, &c1))).unwrap();

        let h1_v = h1.to_vec().unwrap(); // test-only readback
        let h2_v = h2.to_vec().unwrap(); // test-only readback
        let diff: f32 = h1_v.iter().zip(h2_v.iter()).map(|(a, b)| (a - b).abs()).sum();
        assert!(diff > 0.0, "Sequential outputs should differ");
    }
}
