//! 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 std::sync::Arc;

use cudarc::cublas::CudaBlas;
use cudarc::driver::CudaStream;
use serde::{Deserialize, Serialize};

use ml_core::cuda_autograd::{GpuLinear, GpuTensor, GpuVarStore};
use ml_core::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::ConfigError(
                "embed_dim must be greater than 0".to_owned(),
            ));
        }
        if num_heads == 0 {
            return Err(MLError::ConfigError(
                "num_heads must be greater than 0".to_owned(),
            ));
        }
        if embed_dim % num_heads != 0 {
            return Err(MLError::ConfigError(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 const 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,
        }
    }
}

/// Multi-Head Self-Attention Layer
///
/// Implements scaled dot-product attention with multiple heads for
/// capturing different aspects of temporal patterns in the input sequence.
///
/// NOTE: The forward pass operates on host data (CPU) for the attention
/// computation since the GpuLinear/GpuTensor abstraction does not yet support
/// 3D tensor reshaping required for multi-head attention. The hot-path
/// attention in the DQN trainer uses fused CUDA kernels directly.
#[allow(missing_debug_implementations)]
pub struct MultiHeadAttention {
    /// Configuration
    config: MultiHeadAttentionConfig,
    /// Query projection
    wq: GpuLinear,
    /// Key projection
    wk: GpuLinear,
    /// Value projection
    wv: GpuLinear,
    /// Output projection
    wo: GpuLinear,
    /// Variable store holding all parameters
    store: GpuVarStore,
    /// cuBLAS handle
    cublas: CudaBlas,
    /// Layer normalization weights (optional, stored as host vecs for cold path)
    ln_weight: Option<Vec<f32>>,
    ln_bias: Option<Vec<f32>>,
    /// CUDA stream
    stream: Arc<CudaStream>,
}

impl MultiHeadAttention {
    /// Create a new Multi-Head Attention layer
    pub fn new(
        config: MultiHeadAttentionConfig,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let embed_dim = config.embed_dim;

        let mut store = GpuVarStore::new(stream.clone());

        let wq = store.linear("wq", embed_dim, embed_dim)?;
        let wk = store.linear("wk", embed_dim, embed_dim)?;
        let wv = store.linear("wv", embed_dim, embed_dim)?;
        let wo = store.linear("wo", embed_dim, embed_dim)?;

        let cublas = CudaBlas::new(stream.clone()).map_err(|e| {
            MLError::ModelError(format!("cuBLAS init: {e}"))
        })?;

        // Layer normalization parameters (host-side for cold path)
        let (ln_weight, ln_bias) = if config.use_layer_norm {
            (
                Some(vec![1.0_f32; embed_dim]),
                Some(vec![0.0_f32; embed_dim]),
            )
        } else {
            (None, None)
        };

        Ok(Self {
            config,
            wq,
            wk,
            wv,
            wo,
            store,
            cublas,
            ln_weight,
            ln_bias,
            stream: stream.clone(),
        })
    }

    /// Forward pass through the attention layer (cold path, CPU matmul).
    ///
    /// Operates on host-side data. For GPU-resident attention in the training
    /// loop, use the fused CUDA attention kernel directly.
    ///
    /// # Arguments
    ///
    /// * `x_host` - Input data, flat `[batch_size * seq_len * embed_dim]` (host)
    /// * `batch_size` - Batch size
    /// * `seq_len` - Sequence length
    /// * `mask_host` - Optional attention mask `[seq_len * seq_len]` (host)
    ///   Values should be 0 for positions to attend and -inf for positions to mask
    ///
    /// # Returns
    ///
    /// Output data, flat `[batch_size * seq_len * embed_dim]` (host)
    pub fn forward(
        &self,
        x_host: &[f32],
        batch_size: usize,
        seq_len: usize,
        mask_host: Option<&[f32]>,
    ) -> Result<Vec<f32>, MLError> {
        let embed_dim = self.config.embed_dim;
        let expected_len = batch_size * seq_len * embed_dim;
        if x_host.len() != expected_len {
            return Err(MLError::DimensionMismatch {
                expected: expected_len,
                actual: x_host.len(),
            });
        }

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

        // Flatten x into [batch_size * seq_len, embed_dim] for linear projection
        let x_2d = GpuTensor::from_host(
            x_host,
            vec![batch_size * seq_len, embed_dim],
            &self.stream,
        )?;

        // 1. Linear projections (GPU)
        let (q_gpu, _) = self.wq.forward(&x_2d, &self.store, &self.cublas, &self.stream)?;
        let (k_gpu, _) = self.wk.forward(&x_2d, &self.store, &self.cublas, &self.stream)?;
        let (v_gpu, _) = self.wv.forward(&x_2d, &self.store, &self.cublas, &self.stream)?;

        // Download to host for attention computation
        let q = q_gpu.to_host(&self.stream)?;
        let k = k_gpu.to_host(&self.stream)?;
        let v = v_gpu.to_host(&self.stream)?;

        // 2. Reshape to [batch, num_heads, seq_len, head_dim] and compute attention (CPU)
        let scale = (head_dim as f64).sqrt();
        let mut attn_output = vec![0.0_f32; batch_size * seq_len * embed_dim];

        for b in 0..batch_size {
            for h in 0..num_heads {
                // Extract Q, K, V for this batch/head
                // Input shape is [batch*seq_len, embed_dim]
                // Q[b, s, h, d] = q[(b*seq_len + s) * embed_dim + h*head_dim + d]
                let mut scores = vec![0.0_f32; seq_len * seq_len];

                // Compute QK^T / sqrt(d_k)
                for i in 0..seq_len {
                    for j in 0..seq_len {
                        let mut dot = 0.0_f64;
                        for d in 0..head_dim {
                            let qi = q.get((b * seq_len + i) * embed_dim + h * head_dim + d)
                                .copied().unwrap_or(0.0) as f64;
                            let kj = k.get((b * seq_len + j) * embed_dim + h * head_dim + d)
                                .copied().unwrap_or(0.0) as f64;
                            dot += qi * kj;
                        }
                        if let Some(s) = scores.get_mut(i * seq_len + j) {
                            *s = (dot / scale) as f32;
                        }
                    }
                }

                // Apply mask if provided
                if let Some(mask) = mask_host {
                    for i in 0..seq_len {
                        for j in 0..seq_len {
                            let mask_val = mask.get(i * seq_len + j).copied().unwrap_or(0.0);
                            if let Some(s) = scores.get_mut(i * seq_len + j) {
                                *s += mask_val;
                            }
                        }
                    }
                }

                // Softmax per row
                for i in 0..seq_len {
                    let row_start = i * seq_len;
                    let mut max_val = f32::NEG_INFINITY;
                    for j in 0..seq_len {
                        let val = scores.get(row_start + j).copied().unwrap_or(f32::NEG_INFINITY);
                        if val > max_val { max_val = val; }
                    }
                    let mut exp_sum = 0.0_f32;
                    for j in 0..seq_len {
                        if let Some(s) = scores.get_mut(row_start + j) {
                            *s = (*s - max_val).exp();
                            exp_sum += *s;
                        }
                    }
                    for j in 0..seq_len {
                        if let Some(s) = scores.get_mut(row_start + j) {
                            *s /= exp_sum;
                        }
                    }
                }

                // Weighted sum: output = softmax(QK^T/sqrt(d_k)) * V
                for i in 0..seq_len {
                    for d in 0..head_dim {
                        let mut val = 0.0_f32;
                        for j in 0..seq_len {
                            let attn_w = scores.get(i * seq_len + j).copied().unwrap_or(0.0);
                            let vj = v.get((b * seq_len + j) * embed_dim + h * head_dim + d)
                                .copied().unwrap_or(0.0);
                            val += attn_w * vj;
                        }
                        let out_idx = (b * seq_len + i) * embed_dim + h * head_dim + d;
                        if let Some(o) = attn_output.get_mut(out_idx) {
                            *o = val;
                        }
                    }
                }
            }
        }

        // 3. Output projection (GPU)
        let attn_2d = GpuTensor::from_host(
            &attn_output,
            vec![batch_size * seq_len, embed_dim],
            &self.stream,
        )?;
        let (out_gpu, _) = self.wo.forward(&attn_2d, &self.store, &self.cublas, &self.stream)?;
        let mut output = out_gpu.to_host(&self.stream)?;

        // 4. Residual connection
        if self.config.use_residual {
            for (o, x) in output.iter_mut().zip(x_host.iter()) {
                *o += x;
            }
        }

        // 5. Layer normalization (CPU cold path)
        if let (Some(weight), Some(bias)) = (&self.ln_weight, &self.ln_bias) {
            let eps = self.config.layer_norm_eps as f32;
            for sample in 0..(batch_size * seq_len) {
                let base = sample * embed_dim;
                let slice = output.get(base..base + embed_dim).ok_or_else(|| {
                    MLError::ModelError("LN slice out of bounds".into())
                })?;

                // Compute mean and variance
                let mean: f32 = slice.iter().sum::<f32>() / embed_dim as f32;
                let var: f32 = slice.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / embed_dim as f32;
                let inv_std = 1.0 / (var + eps).sqrt();

                // Normalize and apply affine
                for d in 0..embed_dim {
                    if let Some(o) = output.get_mut(base + d) {
                        let w = weight.get(d).copied().unwrap_or(1.0);
                        let b = bias.get(d).copied().unwrap_or(0.0);
                        *o = (*o - mean) * inv_std * w + b;
                    }
                }
            }
        }

        Ok(output)
    }

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

#[cfg(test)]
#[allow(clippy::assertions_on_result_states)]
mod tests {
    use super::*;

    fn make_stream() -> Arc<CudaStream> {
        let device = ml_core::device::MlDevice::cuda(0).expect("CUDA required");
        device.cuda_stream().expect("stream").clone()
    }

    #[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 stream = make_stream();
        let config = MultiHeadAttentionConfig::new(64, 4)?;

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

        Ok(())
    }

    #[test]
    fn test_forward_pass_shape() -> Result<(), MLError> {
        let stream = make_stream();
        let config = MultiHeadAttentionConfig::new(64, 4)?;

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

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

        let input_data = vec![0.1_f32; batch_size * seq_len * embed_dim];

        let output = attention.forward(&input_data, batch_size, seq_len, None)?;

        assert_eq!(output.len(), batch_size * seq_len * embed_dim);

        Ok(())
    }

    #[test]
    fn test_forward_with_mask() -> Result<(), MLError> {
        let stream = make_stream();
        let config = MultiHeadAttentionConfig::new(64, 4)?;

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

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

        let input_data = vec![0.1_f32; batch_size * seq_len * embed_dim];

        // 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 output = attention.forward(&input_data, batch_size, seq_len, Some(&mask_data))?;

        assert_eq!(output.len(), batch_size * seq_len * embed_dim);

        Ok(())
    }

    #[test]
    fn test_dimension_mismatch() -> Result<(), MLError> {
        let stream = make_stream();
        let config = MultiHeadAttentionConfig::new(64, 4)?;

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

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

        let input_data = vec![0.1_f32; batch_size * seq_len * wrong_embed_dim];

        let result = attention.forward(&input_data, batch_size, seq_len, None);
        assert!(result.is_err());

        Ok(())
    }

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

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

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

        let input_data = vec![1.0_f32; batch_size * seq_len * embed_dim];

        let output = attention.forward(&input_data, batch_size, seq_len, None)?;

        // Output should exist and have correct length
        assert_eq!(output.len(), batch_size * seq_len * embed_dim);

        Ok(())
    }

    #[test]
    fn test_multiple_heads() -> Result<(), MLError> {
        let stream = make_stream();

        for num_heads in [1, 2, 4, 8] {
            let embed_dim = 64;
            let config = MultiHeadAttentionConfig::new(embed_dim, num_heads)?;

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

            let batch_size = 2;
            let seq_len = 8;

            let input_data = vec![0.1_f32; batch_size * seq_len * embed_dim];

            let output = attention.forward(&input_data, batch_size, seq_len, None)?;

            assert_eq!(output.len(), batch_size * seq_len * embed_dim);
        }

        Ok(())
    }
}
