//! Simple attention implementation optimized for compilation
//!
//! This module provides basic attention mechanisms using modern Candle API patterns.

use ml_core::device::MlDevice;
use ml_core::native_types::NativeDevice;
use ml_core::cuda_autograd::GpuTensor;

/// A causal attention mask for transformer self-attention.
///
/// The mask is a 2-D `[seq_len, seq_len]` tensor where positions above the
/// diagonal are `f32::NEG_INFINITY` (masked out) and positions on or below
/// the diagonal are `0.0` (allowed). Adding this to raw attention logits
/// prevents the model from attending to future positions.
#[derive(Debug)]
pub struct AttentionMask {
    pub mask: GpuTensor,
}

impl AttentionMask {
    /// Creates a lower-triangular causal mask of shape `[seq_len, seq_len]`.
    ///
    /// Positions `(i, j)` where `j <= i` are `0.0`; positions where `j > i`
    /// are `f32::NEG_INFINITY`.
    pub fn causal(seq_len: usize, device: &NativeDevice) -> Result<Self, ml_core::MLError> {
        let mut data = Vec::with_capacity(seq_len * seq_len);
        for i in 0..seq_len {
            for j in 0..seq_len {
                if j <= i {
                    data.push(0.0_f32);
                } else {
                    data.push(f32::NEG_INFINITY);
                }
            }
        }
        let ordinal = device.cuda_ordinal().ok_or_else(|| {
            ml_core::MLError::DeviceError("Causal mask requires CUDA device".to_owned())
        })?;
        let ml_dev = MlDevice::cuda(ordinal)?;
        let stream = ml_dev.cuda_stream()?;
        let mask = GpuTensor::from_host(&data, vec![seq_len, seq_len], stream)?;
        Ok(Self { mask })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tft::temporal_attention::AttentionConfig;

    #[test]
    fn test_attention_config() {
        let config = AttentionConfig::default();
        assert_eq!(config.hidden_dim, 256);
        assert_eq!(config.num_heads, 8);
        assert_eq!(config.dropout_rate, 0.1);
    }

    #[test]
    fn test_attention_mask() -> Result<(), ml_core::MLError> {
        let device = NativeDevice::Cuda(0);
        let mask = AttentionMask::causal(4, &device)?;
        assert_eq!(mask.mask.dims(), &[4, 4]);

        // Download mask to host for verification
        let ml_dev = MlDevice::cuda(0)?;
        let stream = ml_dev.cuda_stream()?;
        let flat = mask.mask.to_host(stream)?;
        // Row 0: [0, -inf, -inf, -inf]
        assert_eq!(*flat.get(0).unwrap_or(&f32::NAN), 0.0);
        assert!(flat.get(1).unwrap_or(&0.0).is_infinite());
        // Row 1: [0, 0, -inf, -inf]
        assert_eq!(*flat.get(4).unwrap_or(&f32::NAN), 0.0);
        assert_eq!(*flat.get(5).unwrap_or(&f32::NAN), 0.0);
        assert!(flat.get(6).unwrap_or(&0.0).is_infinite());
        // Row 3 (last): [0, 0, 0, 0]
        assert_eq!(*flat.get(12).unwrap_or(&f32::NAN), 0.0);
        assert_eq!(*flat.get(15).unwrap_or(&f32::NAN), 0.0);
        Ok(())
    }
}
