//! mLSTM cell -- matrix memory LSTM with covariance-based updates.
//!
//! Instead of a scalar cell state, mLSTM uses a matrix memory `C_t` of shape
//! `(head_dim, head_dim)` per head. The update rule:
//!   `C_t` = `f_t` * C_{t-1} + `i_t` * (`v_t` ⊗ `k_t^T`)
//!   `h_t` = `o_t` * (`C_t` * `q_t`) / `max(|n_t^T` * `q_t`|, 1)
//!
//! This gives higher memory capacity than scalar LSTM, enabling the model to
//! store and retrieve associations between keys and values.
//!
//! **OOM safeguard:** Matrix memory is `(num_heads, head_dim, head_dim)` per
//! batch item. With default hidden=64, heads=4: `head_dim=16`, so C is 4x16x16=1024
//! floats per sample -- very manageable. Be cautious with hidden>128 or heads<2.

use std::sync::Arc;

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

use crate::gpu_tensor::{
    gpu_add, gpu_cat_dim0, gpu_cat_dim1, gpu_clamp, gpu_exp, gpu_matmul, gpu_mul,
    gpu_narrow_2d, gpu_sigmoid, gpu_tanh, gpu_transpose,
    GpuLinear, GpuTensor,
};

/// mLSTM cell with matrix memory.
#[derive(Debug)]
pub struct MLSTMCell {
    /// Projects input to query, key, value: x -> (q, k, v) each of `hidden_dim`
    qkv_proj: GpuLinear,
    /// Projects hidden state to recurrent qkv
    hidden_proj: GpuLinear,
    /// Gate projections for i, f, o
    gate_proj: GpuLinear,
    hidden_dim: usize,
    num_heads: usize,
    head_dim: usize,
}

impl MLSTMCell {
    /// Create a new mLSTM cell.
    pub fn new(
        input_dim: usize,
        hidden_dim: usize,
        num_heads: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        if hidden_dim % num_heads != 0 {
            return Err(MLError::ConfigError(format!(
                "hidden_dim ({hidden_dim}) must be divisible by num_heads ({num_heads})"
            )));
        }
        let head_dim = hidden_dim / num_heads;

        let qkv_proj = GpuLinear::new(input_dim, 3 * hidden_dim, stream)?;
        let hidden_proj = GpuLinear::new(hidden_dim, 3 * hidden_dim, stream)?;
        let gate_proj = GpuLinear::new(input_dim + hidden_dim, 3 * hidden_dim, stream)?;

        Ok(Self {
            qkv_proj,
            hidden_proj,
            gate_proj,
            hidden_dim,
            num_heads,
            head_dim,
        })
    }

    /// Forward pass through mLSTM cell.
    ///
    /// * `x` - Input tensor `(batch, input_dim)`
    /// * `state` - Optional `(h, c_flat)` where:
    ///   - h: `(batch, hidden_dim)` -- previous hidden state
    ///   - `c_flat`: `(batch, num_heads * head_dim * head_dim)` -- flattened matrix memory
    ///
    /// Returns `(h_new, c_flat_new)`.
    pub fn forward(
        &self,
        x: &GpuTensor,
        state: Option<(&GpuTensor, &GpuTensor)>,
    ) -> Result<(GpuTensor, GpuTensor), MLError> {
        let batch = x.dim(0)?;
        let c_flat_size = self.num_heads * self.head_dim * self.head_dim;

        let (h_prev, c_flat_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, c_flat_size], &x.stream)?;
            (h_zeros, c_zeros)
        };

        // Compute q, k, v from input + hidden recurrence
        let x_qkv = self.qkv_proj.forward(x)?;
        let h_qkv = self.hidden_proj.forward(&h_prev)?;
        let qkv = gpu_add(&x_qkv, &h_qkv)?;

        // Split into q, k, v each of (batch, hidden_dim)
        let q = gpu_narrow_2d(&qkv, 1, 0, self.hidden_dim)?;
        let k = gpu_narrow_2d(&qkv, 1, self.hidden_dim, self.hidden_dim)?;
        let v = gpu_narrow_2d(&qkv, 1, 2 * self.hidden_dim, self.hidden_dim)?;

        // Compute gates from concatenated [x, h_prev]
        let xh = gpu_cat_dim1(x, &h_prev)?;
        let gates_raw = self.gate_proj.forward(&xh)?;

        let i_gate = gpu_narrow_2d(&gates_raw, 1, 0, self.hidden_dim)?;
        let f_gate = gpu_narrow_2d(&gates_raw, 1, self.hidden_dim, self.hidden_dim)?;
        let o_gate = gpu_narrow_2d(&gates_raw, 1, 2 * self.hidden_dim, self.hidden_dim)?;

        // Exponential gates (clamped)
        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)?;
        let o_sig = gpu_sigmoid(&o_gate)?;

        // GPU-native matrix memory update via batched matmul.
        //
        // Strategy: reshape (batch, hidden_dim) -> (batch * num_heads, head_dim)
        // so all per-head ops become standard 2D matmul on the flattened batch.

        let hd = self.head_dim;
        let nh = self.num_heads;

        // Reshape q, k, v to (batch * num_heads, head_dim)
        let q_bh = q.reshape(&[batch * nh, hd])?;
        let k_bh = k.reshape(&[batch * nh, hd])?;
        let v_bh = v.reshape(&[batch * nh, hd])?;

        // Compute per-head scalar gates: mean of i_exp and f_exp across head_dim.
        // Reshape (batch, hidden) -> (batch * num_heads, head_dim) then mean.
        let i_bh = i_exp.reshape(&[batch * nh, hd])?;
        let f_bh = f_exp.reshape(&[batch * nh, hd])?;

        // mean_keepdim_last: (batch*nh, hd) -> (batch*nh, 1)
        let i_mean = crate::gpu_tensor::gpu_mean_keepdim_last(&i_bh)?;
        let f_mean = crate::gpu_tensor::gpu_mean_keepdim_last(&f_bh)?;

        // Outer product per head: v_col * k_row^T -> (head_dim, head_dim) per batch*head.
        // But gpu_matmul is 2D only. We process each batch*head element individually.
        // Actually, we can be smarter: v_bh is (batch*nh, hd), k_bh is (batch*nh, hd).
        // We need per-row outer products. Loop over batch*nh (small: typically 2*4=8).
        let bh = batch * nh;

        // Build C_new for each batch*head using GPU matmul.
        // c_flat_prev: (batch, c_flat_size) where c_flat_size = nh * hd * hd
        // Reshape to (batch * nh, hd, hd) via (batch * nh * hd, hd)
        let c_prev_mat = c_flat_prev.reshape(&[bh * hd, hd])?;

        // For each b*h element: C_new[b*h] = f_mean[b*h] * C_prev[b*h] + i_mean[b*h] * v[b*h] @ k[b*h]^T
        // We can do this without looping by treating (bh * hd, hd) as blocks:
        //
        // Scale entire c_prev by f_mean (broadcast column vector across hd*hd block).
        // We need f_mean_expanded: repeat each scalar hd times -> (bh * hd, 1)
        let f_mean_exp = crate::gpu_tensor::gpu_unsqueeze_broadcast(&f_mean.reshape(&[bh])?, hd)?;
        // f_mean_exp: (bh, hd) -- reshape to (bh * hd, 1) for broadcast_mul_col
        let f_mean_col = f_mean_exp.reshape(&[bh * hd, 1])?;
        let fc = crate::gpu_tensor::gpu_broadcast_mul_col(&c_prev_mat, &f_mean_col)?;

        // Compute v outer k^T scaled by i_mean per head.
        // Loop over bh since gpu_matmul is 2D only. bh is small (batch * num_heads).
        let mut vk_blocks = Vec::with_capacity(bh);
        for idx in 0..bh {
            let v_row = gpu_narrow_2d(&v_bh, 0, idx, 1)?; // (1, hd)
            let k_row = gpu_narrow_2d(&k_bh, 0, idx, 1)?; // (1, hd)
            // v * k^T = (hd,1) @ (1,hd) = (hd, hd)
            let v_col = gpu_transpose(&v_row)?;             // (hd, 1)
            let outer_vk = gpu_matmul(&v_col, &k_row)?;    // (hd, hd)

            // Scale by i_mean[idx]: broadcast (1,1) -> (hd, 1) then col-multiply
            let i_row = gpu_narrow_2d(&i_mean, 0, idx, 1)?; // (1, 1)
            let i_scalar_col = crate::gpu_tensor::gpu_unsqueeze_broadcast(
                &i_row.reshape(&[1])?, hd,
            )?.reshape(&[hd, 1])?;
            let scaled_outer = crate::gpu_tensor::gpu_broadcast_mul_col(&outer_vk, &i_scalar_col)?;
            vk_blocks.push(scaled_outer);
        }
        // Stack vk_blocks into (bh * hd, hd)
        let vk_stacked = gpu_cat_dim0(&vk_blocks, &x.stream)?;

        // C_new = fc + vk_stacked, both (bh * hd, hd)
        let c_new_mat = gpu_add(&fc, &vk_stacked)?;

        // Now compute h_head = C_new * q for each batch*head.
        // C_new_mat: (bh * hd, hd), q_bh: (bh, hd)
        // Per-head: h[b*h] = C_new[b*h] @ q[b*h]^T -> (hd, 1)
        let mut h_blocks = Vec::with_capacity(bh);
        for idx in 0..bh {
            let c_block = gpu_narrow_2d(&c_new_mat, 0, idx * hd, hd)?; // (hd, hd)
            let q_row = gpu_narrow_2d(&q_bh, 0, idx, 1)?;               // (1, hd)
            let q_col = gpu_transpose(&q_row)?;                          // (hd, 1)
            let h_col = gpu_matmul(&c_block, &q_col)?;                   // (hd, 1)
            let h_row = gpu_transpose(&h_col)?;                          // (1, hd)
            h_blocks.push(h_row);
        }
        // Stack: (bh, hd)
        let h_flat_bh = gpu_cat_dim0(&h_blocks, &x.stream)?;
        // Reshape back to (batch, hidden_dim)
        let h_flat = h_flat_bh.reshape(&[batch, self.hidden_dim])?;

        // Apply output gate: h = o * tanh(h_flat)
        let h_tanh = gpu_tanh(&h_flat)?;
        let h_new = gpu_mul(&o_sig, &h_tanh)?;

        // Reshape C_new back to (batch, c_flat_size)
        let c_flat_new = c_new_mat.reshape(&[batch, c_flat_size])?;

        Ok((h_new, c_flat_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_mlstm_output_shape() {
        let stream = test_stream();
        let cell = MLSTMCell::new(8, 16, 2, &stream).unwrap();

        let x = GpuTensor::randn(&[2, 8], 1.0, &stream).unwrap();
        let (h, c) = cell.forward(&x, None).unwrap();
        assert_eq!(h.shape, vec![2, 16]);
        // c_flat: 2 heads * 8 * 8 = 128
        assert_eq!(c.shape, vec![2, 128]);
    }

    #[test]
    fn test_mlstm_sequential() {
        let stream = test_stream();
        let cell = MLSTMCell::new(8, 16, 2, &stream).unwrap();

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

        let x2 = GpuTensor::randn(&[2, 8], 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");
    }

    #[test]
    fn test_mlstm_invalid_head_config() {
        let stream = test_stream();
        let result = MLSTMCell::new(8, 15, 4, &stream); // 15 not divisible by 4
        assert!(result.is_err());
    }
}
