foxhunt/ml/src/dqn/attention.rs

//! Multi-Head Self-Attention for Temporal Pattern Recognition
//!
//! This module implements scaled dot-product attention with multiple heads
//! for capturing temporal dependencies in time-series data. It enables the
//! network to focus on different aspects of the input sequence simultaneously.
//!
//! ## Features
//!
//! - Multi-head attention with configurable heads and dimensions
//! - Scaled dot-product attention with optional masking
//! - Xavier initialization for stable training
//! - Layer normalization for improved training dynamics
//! - Optional residual connections
//!
//! ## Architecture
//!
//! ```text
//! Input (batch, seq_len, embed_dim)
//!   |
//!   ├─> Query (WQ) ─┐
//!   ├─> Key (WK) ───┤
//!   └─> Value (WV) ─┴─> Scaled Dot-Product Attention
//!                         |
//!                         v
//!                    Multi-Head Concat
//!                         |
//!                         v
//!                    Output Linear (WO)
//!                         |
//!                         v
//!                    Output (batch, seq_len, embed_dim)
//! ```

use candle_core::{Device, Result as CandleResult, Tensor};
use candle_nn::{Linear, Module, VarBuilder};
use serde::{Deserialize, Serialize};

use crate::cuda_compat::layer_norm_with_fallback;
use crate::dqn::xavier_init::linear_xavier;
use crate::MLError;

/// Configuration for Multi-Head Attention layer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MultiHeadAttentionConfig {
    /// Embedding dimension (must be divisible by num_heads)
    pub embed_dim: usize,
    /// Number of attention heads
    pub num_heads: usize,
    /// Dropout probability (0.0 to 1.0)
    pub dropout: f64,
    /// Whether to use layer normalization
    pub use_layer_norm: bool,
    /// LayerNorm epsilon for numerical stability
    pub layer_norm_eps: f64,
    /// Whether to use residual connections
    pub use_residual: bool,
}

impl MultiHeadAttentionConfig {
    /// Create a new configuration with validation
    pub fn new(embed_dim: usize, num_heads: usize) -> Result<Self, MLError> {
        if embed_dim == 0 {
            return Err(MLError::ConfigurationError(
                "embed_dim must be greater than 0".to_string(),
            ));
        }
        if num_heads == 0 {
            return Err(MLError::ConfigurationError(
                "num_heads must be greater than 0".to_string(),
            ));
        }
        if embed_dim % num_heads != 0 {
            return Err(MLError::ConfigurationError(format!(
                "embed_dim ({}) must be divisible by num_heads ({})",
                embed_dim, num_heads
            )));
        }

        Ok(Self {
            embed_dim,
            num_heads,
            dropout: 0.1,
            use_layer_norm: true,
            layer_norm_eps: 1e-5,
            use_residual: true,
        })
    }

    /// Get the dimension of each attention head
    pub fn head_dim(&self) -> usize {
        self.embed_dim / self.num_heads
    }
}

impl Default for MultiHeadAttentionConfig {
    fn default() -> Self {
        Self {
            embed_dim: 64,
            num_heads: 4,
            dropout: 0.1,
            use_layer_norm: true,
            layer_norm_eps: 1e-5,
            use_residual: true,
        }
    }
}

/// LayerNorm parameters for attention layer
#[derive(Debug)]
struct AttentionLayerNorm {
    weight: Tensor,
    bias: Tensor,
    normalized_shape: usize,
    eps: f64,
}

impl AttentionLayerNorm {
    fn new(
        normalized_shape: usize,
        eps: f64,
        var_builder: &VarBuilder<'_>,
        name: &str,
    ) -> CandleResult<Self> {
        let weight = var_builder.get(normalized_shape, &format!("{}_weight", name))?;
        let bias = var_builder.get(normalized_shape, &format!("{}_bias", name))?;

        Ok(Self {
            weight,
            bias,
            normalized_shape,
            eps,
        })
    }

    fn forward(&self, x: &Tensor) -> CandleResult<Tensor> {
        layer_norm_with_fallback(
            x,
            &[self.normalized_shape],
            Some(&self.weight),
            Some(&self.bias),
            self.eps,
        )
        .map_err(|e| candle_core::Error::Msg(format!("LayerNorm failed: {}", e)))
    }
}

/// Multi-Head Self-Attention Layer
///
/// Implements scaled dot-product attention with multiple heads for
/// capturing different aspects of temporal patterns in the input sequence.
#[allow(missing_debug_implementations)]
pub struct MultiHeadAttention {
    /// Configuration
    config: MultiHeadAttentionConfig,
    /// Query projection
    wq: Linear,
    /// Key projection
    wk: Linear,
    /// Value projection
    wv: Linear,
    /// Output projection
    wo: Linear,
    /// Layer normalization (optional)
    layer_norm: Option<AttentionLayerNorm>,
    /// Compute device
    device: Device,
}

impl MultiHeadAttention {
    /// Create a new Multi-Head Attention layer
    ///
    /// # Arguments
    ///
    /// * `config` - Attention configuration
    /// * `var_builder` - Variable builder for parameter initialization
    /// * `device` - Compute device (CPU/CUDA)
    ///
    /// # Returns
    ///
    /// Returns `Ok(MultiHeadAttention)` on success, or error if initialization fails
    pub fn new(
        config: MultiHeadAttentionConfig,
        var_builder: &VarBuilder<'_>,
        device: &Device,
    ) -> Result<Self, MLError> {
        let embed_dim = config.embed_dim;

        // Create Q, K, V projections with Xavier initialization
        let wq = linear_xavier(embed_dim, embed_dim, var_builder.pp("wq"))
            .map_err(|e| MLError::InitializationError {
                component: "wq".to_string(),
                message: e.to_string(),
            })?;

        let wk = linear_xavier(embed_dim, embed_dim, var_builder.pp("wk"))
            .map_err(|e| MLError::InitializationError {
                component: "wk".to_string(),
                message: e.to_string(),
            })?;

        let wv = linear_xavier(embed_dim, embed_dim, var_builder.pp("wv"))
            .map_err(|e| MLError::InitializationError {
                component: "wv".to_string(),
                message: e.to_string(),
            })?;

        // Output projection
        let wo = linear_xavier(embed_dim, embed_dim, var_builder.pp("wo"))
            .map_err(|e| MLError::InitializationError {
                component: "wo".to_string(),
                message: e.to_string(),
            })?;

        // Optional layer normalization
        let layer_norm = if config.use_layer_norm {
            Some(
                AttentionLayerNorm::new(embed_dim, config.layer_norm_eps, var_builder, "ln")
                    .map_err(|e| MLError::InitializationError {
                        component: "layer_norm".to_string(),
                        message: e.to_string(),
                    })?,
            )
        } else {
            None
        };

        Ok(Self {
            config,
            wq,
            wk,
            wv,
            wo,
            layer_norm,
            device: device.clone(),
        })
    }

    /// Forward pass through the attention layer
    ///
    /// # Arguments
    ///
    /// * `x` - Input tensor of shape `(batch_size, seq_len, embed_dim)`
    /// * `mask` - Optional attention mask of shape `(batch_size, seq_len, seq_len)` or `(seq_len, seq_len)`
    ///            Values should be 0 for positions to attend and -inf for positions to mask
    ///
    /// # Returns
    ///
    /// Returns tensor of shape `(batch_size, seq_len, embed_dim)`
    ///
    /// # Algorithm
    ///
    /// 1. Linear projections: Q = XW_Q, K = XW_K, V = XW_V
    /// 2. Split into multiple heads
    /// 3. Scaled dot-product attention: Attention(Q,K,V) = softmax(QK^T/√d_k)V
    /// 4. Concatenate heads and apply output projection
    /// 5. Optional: Add residual connection and layer normalization
    pub fn forward(&self, x: &Tensor, mask: Option<&Tensor>) -> Result<Tensor, MLError> {
        let residual = x.clone();

        // Get dimensions
        let (batch_size, seq_len, embed_dim) = x
            .dims3()
            .map_err(|e| MLError::InvalidInput(format!("Expected 3D input tensor: {}", e)))?;

        if embed_dim != self.config.embed_dim {
            return Err(MLError::DimensionMismatch {
                expected: self.config.embed_dim,
                actual: embed_dim,
            });
        }

        let num_heads = self.config.num_heads;
        let head_dim = self.config.head_dim();

        // 1. Linear projections
        let q = self
            .wq
            .forward(x)
            .map_err(|e| MLError::ModelError(format!("Query projection failed: {}", e)))?;
        let k = self
            .wk
            .forward(x)
            .map_err(|e| MLError::ModelError(format!("Key projection failed: {}", e)))?;
        let v = self
            .wv
            .forward(x)
            .map_err(|e| MLError::ModelError(format!("Value projection failed: {}", e)))?;

        // 2. Reshape for multi-head attention: (batch, seq_len, embed_dim) -> (batch, num_heads, seq_len, head_dim)
        let q = self.reshape_for_attention(&q, batch_size, seq_len, num_heads, head_dim)?;
        let k = self.reshape_for_attention(&k, batch_size, seq_len, num_heads, head_dim)?;
        let v = self.reshape_for_attention(&v, batch_size, seq_len, num_heads, head_dim)?;

        // 3. Scaled dot-product attention
        let attn_output = self.scaled_dot_product_attention(&q, &k, &v, mask, head_dim)?;

        // 4. Reshape back: (batch, num_heads, seq_len, head_dim) -> (batch, seq_len, embed_dim)
        let attn_output = attn_output
            .transpose(1, 2)
            .map_err(|e| MLError::TensorOperationError(format!("Transpose failed: {}", e)))?
            .reshape((batch_size, seq_len, embed_dim))
            .map_err(|e| MLError::TensorOperationError(format!("Reshape failed: {}", e)))?;

        // 5. Output projection
        let mut output = self
            .wo
            .forward(&attn_output)
            .map_err(|e| MLError::ModelError(format!("Output projection failed: {}", e)))?;

        // 6. Residual connection
        if self.config.use_residual {
            output = (output + residual).map_err(|e| {
                MLError::TensorOperationError(format!("Residual connection failed: {}", e))
            })?;
        }

        // 7. Layer normalization
        if let Some(ref ln) = self.layer_norm {
            output = ln
                .forward(&output)
                .map_err(|e| MLError::ModelError(format!("Layer normalization failed: {}", e)))?;
        }

        Ok(output)
    }

    /// Reshape tensor for multi-head attention
    fn reshape_for_attention(
        &self,
        x: &Tensor,
        batch_size: usize,
        seq_len: usize,
        num_heads: usize,
        head_dim: usize,
    ) -> Result<Tensor, MLError> {
        x.reshape((batch_size, seq_len, num_heads, head_dim))
            .map_err(|e| MLError::TensorOperationError(format!("Reshape failed: {}", e)))?
            .transpose(1, 2)
            .map_err(|e| MLError::TensorOperationError(format!("Transpose failed: {}", e)))
    }

    /// Scaled dot-product attention
    ///
    /// Attention(Q, K, V) = softmax(QK^T / √d_k) V
    fn scaled_dot_product_attention(
        &self,
        q: &Tensor,
        k: &Tensor,
        v: &Tensor,
        mask: Option<&Tensor>,
        head_dim: usize,
    ) -> Result<Tensor, MLError> {
        // QK^T
        let k_transposed = k
            .transpose(2, 3)
            .map_err(|e| MLError::TensorOperationError(format!("Key transpose failed: {}", e)))?;

        let mut scores = q
            .matmul(&k_transposed)
            .map_err(|e| MLError::TensorOperationError(format!("QK^T matmul failed: {}", e)))?;

        // Scale by √d_k
        let scale = (head_dim as f64).sqrt();
        scores = (scores / scale)
            .map_err(|e| MLError::TensorOperationError(format!("Scaling failed: {}", e)))?;

        // Apply mask if provided
        if let Some(mask) = mask {
            // Expand mask to match attention scores shape if needed
            let mask_expanded = if mask.rank() == 2 {
                // (seq_len, seq_len) -> (1, 1, seq_len, seq_len)
                mask.unsqueeze(0)
                    .map_err(|e| {
                        MLError::TensorOperationError(format!("Mask unsqueeze failed: {}", e))
                    })?
                    .unsqueeze(0)
                    .map_err(|e| {
                        MLError::TensorOperationError(format!("Mask unsqueeze failed: {}", e))
                    })?
            } else {
                mask.clone()
            };

            scores = (scores + mask_expanded).map_err(|e| {
                MLError::TensorOperationError(format!("Mask application failed: {}", e))
            })?;
        }

        // Softmax over the last dimension
        let attn_weights = candle_nn::ops::softmax(&scores, 3).map_err(|e| {
            MLError::TensorOperationError(format!("Softmax failed: {}", e))
        })?;

        // Apply attention weights to values
        attn_weights
            .matmul(v)
            .map_err(|e| MLError::TensorOperationError(format!("Attention matmul failed: {}", e)))
    }

    /// Get the configuration
    pub fn config(&self) -> &MultiHeadAttentionConfig {
        &self.config
    }

    /// Get the device
    pub fn device(&self) -> &Device {
        &self.device
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_nn::VarMap;

    #[test]
    fn test_config_validation() {
        // Valid config
        let config = MultiHeadAttentionConfig::new(64, 4);
        assert!(config.is_ok());
        let config = config.expect("Should succeed");
        assert_eq!(config.head_dim(), 16);

        // Invalid: embed_dim = 0
        let config = MultiHeadAttentionConfig::new(0, 4);
        assert!(config.is_err());

        // Invalid: num_heads = 0
        let config = MultiHeadAttentionConfig::new(64, 0);
        assert!(config.is_err());

        // Invalid: embed_dim not divisible by num_heads
        let config = MultiHeadAttentionConfig::new(65, 4);
        assert!(config.is_err());
    }

    #[test]
    fn test_default_config() {
        let config = MultiHeadAttentionConfig::default();
        assert_eq!(config.embed_dim, 64);
        assert_eq!(config.num_heads, 4);
        assert_eq!(config.head_dim(), 16);
        assert!(config.use_layer_norm);
        assert!(config.use_residual);
    }

    #[test]
    fn test_attention_creation() -> Result<(), MLError> {
        let device = Device::Cpu;
        let config = MultiHeadAttentionConfig::new(64, 4)?;
        let vars = VarMap::new();
        let vb = VarBuilder::from_varmap(&vars, DType::F32, &device);

        let attention = MultiHeadAttention::new(config, &vb, &device)?;
        assert_eq!(attention.config().embed_dim, 64);
        assert_eq!(attention.config().num_heads, 4);

        Ok(())
    }

    #[test]
    fn test_forward_pass_shape() -> Result<(), MLError> {
        let device = Device::Cpu;
        let config = MultiHeadAttentionConfig::new(64, 4)?;
        let vars = VarMap::new();
        let vb = VarBuilder::from_varmap(&vars, DType::F32, &device);

        let attention = MultiHeadAttention::new(config, &vb, &device)?;

        // Create input: (batch=2, seq_len=8, embed_dim=64)
        let batch_size = 2;
        let seq_len = 8;
        let embed_dim = 64;

        let input_data = vec![0.1f32; batch_size * seq_len * embed_dim];
        let input = Tensor::from_vec(input_data, (batch_size, seq_len, embed_dim), &device)
            .map_err(|e| MLError::ModelError(format!("Failed to create input: {}", e)))?;

        let output = attention.forward(&input, None)?;

        // Check output shape
        let output_shape = output.dims();
        assert_eq!(output_shape.len(), 3);
        assert_eq!(output_shape[0], batch_size);
        assert_eq!(output_shape[1], seq_len);
        assert_eq!(output_shape[2], embed_dim);

        Ok(())
    }

    #[test]
    fn test_forward_with_mask() -> Result<(), MLError> {
        let device = Device::Cpu;
        let config = MultiHeadAttentionConfig::new(64, 4)?;
        let vars = VarMap::new();
        let vb = VarBuilder::from_varmap(&vars, DType::F32, &device);

        let attention = MultiHeadAttention::new(config, &vb, &device)?;

        // Create input
        let batch_size = 2;
        let seq_len = 8;
        let embed_dim = 64;

        let input_data = vec![0.1f32; batch_size * seq_len * embed_dim];
        let input = Tensor::from_vec(input_data, (batch_size, seq_len, embed_dim), &device)
            .map_err(|e| MLError::ModelError(format!("Failed to create input: {}", e)))?;

        // Create causal mask (lower triangular)
        let mut mask_data = vec![f32::NEG_INFINITY; seq_len * seq_len];
        for i in 0..seq_len {
            for j in 0..=i {
                mask_data[i * seq_len + j] = 0.0;
            }
        }
        let mask = Tensor::from_vec(mask_data, (seq_len, seq_len), &device)
            .map_err(|e| MLError::ModelError(format!("Failed to create mask: {}", e)))?;

        let output = attention.forward(&input, Some(&mask))?;

        // Check output shape
        let output_shape = output.dims();
        assert_eq!(output_shape.len(), 3);
        assert_eq!(output_shape[0], batch_size);
        assert_eq!(output_shape[1], seq_len);
        assert_eq!(output_shape[2], embed_dim);

        Ok(())
    }

    #[test]
    fn test_dimension_mismatch() -> Result<(), MLError> {
        let device = Device::Cpu;
        let config = MultiHeadAttentionConfig::new(64, 4)?;
        let vars = VarMap::new();
        let vb = VarBuilder::from_varmap(&vars, DType::F32, &device);

        let attention = MultiHeadAttention::new(config, &vb, &device)?;

        // Create input with wrong embed_dim
        let batch_size = 2;
        let seq_len = 8;
        let wrong_embed_dim = 32;

        let input_data = vec![0.1f32; batch_size * seq_len * wrong_embed_dim];
        let input = Tensor::from_vec(input_data, (batch_size, seq_len, wrong_embed_dim), &device)
            .map_err(|e| MLError::ModelError(format!("Failed to create input: {}", e)))?;

        let result = attention.forward(&input, None);
        assert!(result.is_err());

        if let Err(MLError::DimensionMismatch { expected, actual }) = result {
            assert_eq!(expected, 64);
            assert_eq!(actual, 32);
        } else {
            panic!("Expected DimensionMismatch error");
        }

        Ok(())
    }

    #[test]
    fn test_residual_connection() -> Result<(), MLError> {
        let device = Device::Cpu;
        let mut config = MultiHeadAttentionConfig::new(64, 4)?;
        config.use_residual = true;
        config.use_layer_norm = false; // Disable to test residual alone

        let vars = VarMap::new();
        let vb = VarBuilder::from_varmap(&vars, DType::F32, &device);

        let attention = MultiHeadAttention::new(config, &vb, &device)?;

        let batch_size = 2;
        let seq_len = 8;
        let embed_dim = 64;

        let input_data = vec![1.0f32; batch_size * seq_len * embed_dim];
        let input = Tensor::from_vec(input_data, (batch_size, seq_len, embed_dim), &device)
            .map_err(|e| MLError::ModelError(format!("Failed to create input: {}", e)))?;

        let output = attention.forward(&input, None)?;

        // Output should exist and have correct shape
        let output_shape = output.dims();
        assert_eq!(output_shape.len(), 3);
        assert_eq!(output_shape[0], batch_size);

        Ok(())
    }

    #[test]
    fn test_multiple_heads() -> Result<(), MLError> {
        let device = Device::Cpu;

        // Test different head configurations
        for num_heads in [1, 2, 4, 8] {
            let embed_dim = 64;
            let config = MultiHeadAttentionConfig::new(embed_dim, num_heads)?;
            let vars = VarMap::new();
            let vb = VarBuilder::from_varmap(&vars, DType::F32, &device);

            let attention = MultiHeadAttention::new(config, &vb, &device)?;

            let batch_size = 2;
            let seq_len = 8;

            let input_data = vec![0.1f32; batch_size * seq_len * embed_dim];
            let input = Tensor::from_vec(input_data, (batch_size, seq_len, embed_dim), &device)
                .map_err(|e| MLError::ModelError(format!("Failed to create input: {}", e)))?;

            let output = attention.forward(&input, None)?;

            // Verify output shape
            let output_shape = output.dims();
            assert_eq!(output_shape[0], batch_size);
            assert_eq!(output_shape[1], seq_len);
            assert_eq!(output_shape[2], embed_dim);
        }

        Ok(())
    }
}