//! 双向多头自注意力机制
//!
//! EmbeddingGemma 使用双向注意力（非因果），这意味着每个 token 都可以关注
//! 序列中的所有其他 token，包括前面和后面的 token。
//!
//! # 与 Causal Attention 的区别
//!
//! - **Causal (LLM)**: 使用下三角掩码，token 只能看到之前的 token
//! - **Bidirectional (Embedding)**: 不使用因果掩码，token 可以看到所有其他 token
//!
//! # 注意力计算
//!
//! ```text
//! Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V
//! ```

use candle_core::{Module, Result, Tensor, D};
use candle_nn::VarBuilder;

use super::config::EmbeddingGemmaConfig;
use super::layers::RMSNorm;
use super::rope::RotaryEmbedding;

/// 双向多头自注意力 (支持 Grouped-Query Attention)
///
/// 实现了 Gemma 风格的多头注意力，但移除了因果掩码以支持双向上下文。
/// 支持 GQA (Grouped-Query Attention) 和 QK Normalization。
pub struct BidirectionalAttention {
    /// Query 投影
    q_proj: candle_nn::Linear,
    /// Key 投影
    k_proj: candle_nn::Linear,
    /// Value 投影
    v_proj: candle_nn::Linear,
    /// 输出投影
    o_proj: candle_nn::Linear,
    /// Query Normalization
    q_norm: RMSNorm,
    /// Key Normalization
    k_norm: RMSNorm,
    /// Query 注意力头数
    num_heads: usize,
    /// Key-Value 注意力头数 (for GQA)
    num_kv_heads: usize,
    /// 每个头的维度
    head_dim: usize,
    /// 缩放因子 (1 / sqrt(query_pre_attn_scalar))
    scale: f64,
    /// 隐藏层维度
    hidden_size: usize,
}

impl BidirectionalAttention {
    /// 创建双向注意力层
    ///
    /// # Arguments
    /// * `config` - 模型配置
    /// * `vb` - 变量构建器
    pub fn new(config: &EmbeddingGemmaConfig, vb: VarBuilder) -> Result<Self> {
        let hidden_size = config.hidden_size;
        let num_heads = config.num_attention_heads;
        let num_kv_heads = config.num_key_value_heads;
        let head_dim = config.head_dim;
        let q_dim = num_heads * head_dim;
        let kv_dim = num_kv_heads * head_dim;

        // Gemma 使用无偏置的线性层
        let q_proj = candle_nn::linear_no_bias(hidden_size, q_dim, vb.pp("q_proj"))?;
        let k_proj = candle_nn::linear_no_bias(hidden_size, kv_dim, vb.pp("k_proj"))?;
        let v_proj = candle_nn::linear_no_bias(hidden_size, kv_dim, vb.pp("v_proj"))?;
        let o_proj = candle_nn::linear_no_bias(q_dim, hidden_size, vb.pp("o_proj"))?;

        // QK Normalization
        let q_norm = RMSNorm::new(head_dim, config.rms_norm_eps, vb.pp("q_norm"))?;
        let k_norm = RMSNorm::new(head_dim, config.rms_norm_eps, vb.pp("k_norm"))?;

        // EmbeddingGemma uses query_pre_attn_scalar for scaling
        let scale = (config.query_pre_attn_scalar as f64).powf(-0.5);

        Ok(Self {
            q_proj,
            k_proj,
            v_proj,
            o_proj,
            q_norm,
            k_norm,
            num_heads,
            num_kv_heads,
            head_dim,
            scale,
            hidden_size,
        })
    }

    /// 前向传播
    ///
    /// # Arguments
    /// * `hidden_states` - 输入张量，形状: (batch, seq_len, hidden_size)
    /// * `attention_mask` - 可选的注意力掩码，形状: (batch, 1, seq_len, seq_len)
    /// * `rope` - 旋转位置编码
    ///
    /// # Returns
    /// 输出张量，形状: (batch, seq_len, hidden_size)
    pub fn forward(
        &self,
        hidden_states: &Tensor,
        attention_mask: Option<&Tensor>,
        rope: &RotaryEmbedding,
    ) -> Result<Tensor> {
        let (batch_size, seq_len, _) = hidden_states.dims3()?;

        // 投影 Q, K, V
        let q = self.q_proj.forward(hidden_states)?;
        let k = self.k_proj.forward(hidden_states)?;
        let v = self.v_proj.forward(hidden_states)?;

        // 重塑为多头格式
        // Q: (batch, seq, num_heads * head_dim) -> (batch, seq, num_heads, head_dim)
        let q = q.reshape((batch_size, seq_len, self.num_heads, self.head_dim))?;
        // K, V: (batch, seq, num_kv_heads * head_dim) -> (batch, seq, num_kv_heads, head_dim)
        let k = k.reshape((batch_size, seq_len, self.num_kv_heads, self.head_dim))?;
        let v = v.reshape((batch_size, seq_len, self.num_kv_heads, self.head_dim))?;

        // 应用 QK Normalization (在每个头的维度上)
        let q = self.q_norm.forward(&q)?;
        let k = self.k_norm.forward(&k)?;

        // 转置为: (batch, num_heads, seq, head_dim) 并确保连续
        let q = q.transpose(1, 2)?.contiguous()?;
        let k = k.transpose(1, 2)?.contiguous()?;
        let v = v.transpose(1, 2)?.contiguous()?;

        // 应用 RoPE
        let (q, k) = rope.apply(&q, &k, 0)?;

        // GQA: 如果 num_kv_heads < num_heads，需要重复 K 和 V
        let (k, v) = if self.num_kv_heads < self.num_heads {
            let n_rep = self.num_heads / self.num_kv_heads;
            // (batch, num_kv_heads, seq, head_dim) -> (batch, num_kv_heads, 1, seq, head_dim)
            let k = k.unsqueeze(2)?;
            let v = v.unsqueeze(2)?;
            // -> (batch, num_kv_heads, n_rep, seq, head_dim)
            let k = k.expand((batch_size, self.num_kv_heads, n_rep, seq_len, self.head_dim))?;
            let v = v.expand((batch_size, self.num_kv_heads, n_rep, seq_len, self.head_dim))?;
            // -> (batch, num_heads, seq, head_dim)
            let k = k.reshape((batch_size, self.num_heads, seq_len, self.head_dim))?;
            let v = v.reshape((batch_size, self.num_heads, seq_len, self.head_dim))?;
            (k.contiguous()?, v.contiguous()?)
        } else {
            (k, v)
        };

        // 计算注意力分数: Q @ K^T * scale
        // (batch, heads, seq, head_dim) @ (batch, heads, head_dim, seq) -> (batch, heads, seq, seq)
        let k_t = k.transpose(D::Minus2, D::Minus1)?.contiguous()?;
        let attn_weights = q.matmul(&k_t)?;
        let attn_weights = (attn_weights * self.scale)?;

        // 应用注意力掩码（如果有）
        // 注意：对于双向注意力，掩码通常只用于 padding
        let attn_weights = match attention_mask {
            Some(mask) => attn_weights.broadcast_add(mask)?,
            None => attn_weights,
        };

        // Softmax
        let attn_weights = candle_nn::ops::softmax(&attn_weights, D::Minus1)?;

        // 注意力输出: attn_weights @ V
        // (batch, heads, seq, seq) @ (batch, heads, seq, head_dim) -> (batch, heads, seq, head_dim)
        let v_cont = v.contiguous()?;
        let attn_output = attn_weights.matmul(&v_cont)?;

        // 转置回: (batch, seq, num_heads, head_dim) 并确保连续
        let attn_output = attn_output.transpose(1, 2)?.contiguous()?;

        // 合并多头: (batch, seq, num_heads * head_dim)
        let attn_output = attn_output.reshape((batch_size, seq_len, self.num_heads * self.head_dim))?;

        // 输出投影
        self.o_proj.forward(&attn_output)
    }

    /// 获取注意力头数
    pub fn num_heads(&self) -> usize {
        self.num_heads
    }

    /// 获取每个头的维度
    pub fn head_dim(&self) -> usize {
        self.head_dim
    }

    /// 获取隐藏层维度
    pub fn hidden_size(&self) -> usize {
        self.hidden_size
    }
}

/// 创建 padding 注意力掩码
///
/// 将 attention_mask (1 表示有效, 0 表示 padding) 转换为
/// 加性掩码 (0 表示有效, -inf 表示 padding)。
///
/// # Arguments
/// * `attention_mask` - 输入掩码，形状: (batch, seq_len)，值为 0 或 1
///
/// # Returns
/// 扩展的掩码，形状: (batch, 1, 1, seq_len)，值为 0 或 -inf
pub fn create_attention_mask(attention_mask: &Tensor) -> Result<Tensor> {
    let dtype = attention_mask.dtype();
    let device = attention_mask.device();

    // 转换: 1 -> 0, 0 -> -inf
    // 使用 where_cond 避免 0 * -inf = NaN 的问题
    let zeros = attention_mask.zeros_like()?;
    let neg_inf = Tensor::full(f32::NEG_INFINITY, attention_mask.shape(), device)?
        .to_dtype(dtype)?;
    
    // 如果 mask == 1，返回 0；否则返回 -inf
    let mask = attention_mask.ge(0.5f64)?.where_cond(&zeros, &neg_inf)?;

    // 扩展维度: (batch, seq) -> (batch, 1, 1, seq)
    let mask = mask.unsqueeze(1)?.unsqueeze(1)?;

    mask.to_dtype(dtype)
}

/// 创建双向注意力的完整掩码（包含 padding）
///
/// 对于双向注意力，我们不需要因果掩码，只需要处理 padding。
///
/// # Arguments
/// * `attention_mask` - padding 掩码，形状: (batch, seq_len)
/// * `seq_len` - 序列长度
///
/// # Returns
/// 完整的注意力掩码，形状: (batch, 1, seq_len, seq_len)
pub fn create_bidirectional_mask(
    attention_mask: &Tensor,
    _seq_len: usize,
) -> Result<Tensor> {
    // 对于双向注意力，我们只需要扩展 padding 掩码
    // 从 (batch, seq) 扩展到 (batch, 1, 1, seq)
    // 然后广播到 (batch, 1, seq, seq)
    create_attention_mask(attention_mask)
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_core::{DType, Device};
    use candle_nn::VarMap;

    fn create_test_varbuilder(device: &Device) -> VarBuilder<'static> {
        let varmap = VarMap::new();
        VarBuilder::from_varmap(&varmap, DType::F32, device)
    }

    #[test]
    fn test_attention_creation() {
        let device = Device::Cpu;
        let vb = create_test_varbuilder(&device);
        let config = EmbeddingGemmaConfig::default();

        let attn = BidirectionalAttention::new(&config, vb);
        assert!(attn.is_ok());

        let attn = attn.unwrap();
        // EmbeddingGemma uses GQA with 3 query heads
        assert_eq!(attn.num_heads(), 3);
        assert_eq!(attn.head_dim(), 256);
        assert_eq!(attn.hidden_size(), 768);
    }

    #[test]
    fn test_attention_forward_shape() {
        let device = Device::Cpu;
        let vb = create_test_varbuilder(&device);
        let config = EmbeddingGemmaConfig::default();

        let attn = BidirectionalAttention::new(&config, vb).unwrap();

        // 创建 RoPE
        let rope = RotaryEmbedding::new(
            config.head_dim,
            config.max_position_embeddings,
            config.rope_theta,
            DType::F32,
            &device,
        )
        .unwrap();

        // 测试输入 (batch=2, seq=16, hidden=768)
        let hidden_states = Tensor::randn(0.0f32, 1.0, (2, 16, 768), &device).unwrap();

        let output = attn.forward(&hidden_states, None, &rope).unwrap();

        // 输出形状应该与输入相同
        assert_eq!(output.dims(), hidden_states.dims());
    }

    #[test]
    fn test_attention_with_mask() {
        let device = Device::Cpu;
        let vb = create_test_varbuilder(&device);
        let config = EmbeddingGemmaConfig::default();

        let attn = BidirectionalAttention::new(&config, vb).unwrap();

        let rope = RotaryEmbedding::new(
            config.head_dim,
            config.max_position_embeddings,
            config.rope_theta,
            DType::F32,
            &device,
        )
        .unwrap();

        // 测试输入
        let batch_size = 2;
        let seq_len = 16;
        let hidden_states =
            Tensor::randn(0.0f32, 1.0, (batch_size, seq_len, 768), &device).unwrap();

        // 创建 padding 掩码 (假设后 4 个 token 是 padding)
        let mask_data: Vec<f32> = (0..batch_size)
            .flat_map(|_| {
                (0..seq_len)
                    .map(|i| if i < seq_len - 4 { 1.0 } else { 0.0 })
                    .collect::<Vec<_>>()
            })
            .collect();
        let attention_mask = Tensor::from_vec(mask_data, (batch_size, seq_len), &device).unwrap();

        // 创建注意力掩码
        let mask = create_attention_mask(&attention_mask).unwrap();

        let output = attn.forward(&hidden_states, Some(&mask), &rope).unwrap();

        assert_eq!(output.dims(), hidden_states.dims());
    }

    #[test]
    fn test_create_attention_mask() {
        let device = Device::Cpu;

        // 创建测试掩码: [1, 1, 1, 0, 0]
        let mask = Tensor::new(&[[1.0f32, 1.0, 1.0, 0.0, 0.0]], &device).unwrap();
        let attn_mask = create_attention_mask(&mask).unwrap();

        // 检查形状: (1, 1, 1, 5)
        assert_eq!(attn_mask.dims(), &[1, 1, 1, 5]);

        // 检查值: 前 3 个应该是 0，后 2 个应该是 -inf
        let values: Vec<f32> = attn_mask.flatten_all().unwrap().to_vec1().unwrap();
        assert_eq!(values[0], 0.0);
        assert_eq!(values[1], 0.0);
        assert_eq!(values[2], 0.0);
        assert!(values[3].is_infinite() && values[3] < 0.0);
        assert!(values[4].is_infinite() && values[4] < 0.0);
    }

    #[test]
    fn test_attention_is_bidirectional() {
        // 验证注意力是双向的（没有因果掩码）
        // 这通过检查前面的 token 可以关注后面的 token 来验证
        let device = Device::Cpu;
        let vb = create_test_varbuilder(&device);
        let config = EmbeddingGemmaConfig {
            num_hidden_layers: 1,
            num_attention_heads: 2,
            head_dim: 32,
            hidden_size: 64,
            intermediate_size: 256,
            ..Default::default()
        };

        let attn = BidirectionalAttention::new(&config, vb).unwrap();

        let rope = RotaryEmbedding::new(
            config.head_dim,
            config.max_position_embeddings,
            config.rope_theta,
            DType::F32,
            &device,
        )
        .unwrap();

        // 创建特殊的输入，使得某些 token 有独特的特征
        let mut data = vec![0.0f32; 1 * 4 * 64];
        // 第一个 token 全为 1
        for i in 0..64 {
            data[i] = 1.0;
        }
        // 最后一个 token 全为 2
        for i in (3 * 64)..(4 * 64) {
            data[i] = 2.0;
        }

        let hidden_states = Tensor::from_vec(data, (1, 4, 64), &device).unwrap();
        let output = attn.forward(&hidden_states, None, &rope).unwrap();

        // 如果是双向的，第一个 token 的输出应该受到最后一个 token 的影响
        // （具体值依赖于随机初始化的权重，这里只验证可以执行）
        assert_eq!(output.dims(), &[1, 4, 64]);
    }
}
