//! Multi-head attention with Grouped-Query Attention (GQA) support.
//!
//! CPU naive implementation. Flash Attention added in F2.

use crate::tensor::Tensor;

/// Attention parameters collected into a struct to avoid excessive function arguments.
pub struct AttentionParams {
    pub num_heads: usize,
    pub num_kv_heads: usize,
    pub head_dim: usize,
    pub seq_len: usize,
    pub kv_len: usize,
    pub causal: bool,
    pub start_pos: usize,
}

/// Compute scaled dot-product attention (CPU, GQA-aware).
///
/// Q: [seq_len, num_heads * head_dim]
/// K: [kv_len, num_kv_heads * head_dim]
/// V: [kv_len, num_kv_heads * head_dim]
///
/// Returns: [seq_len, num_heads * head_dim]
pub fn attention(q: &Tensor, k: &Tensor, v: &Tensor, params: &AttentionParams) -> Tensor {
    let num_heads = params.num_heads;
    let num_kv_heads = params.num_kv_heads;
    let head_dim = params.head_dim;
    let seq_len = params.seq_len;
    let kv_len = params.kv_len;
    let causal = params.causal;
    let start_pos = params.start_pos;
    let heads_per_kv = num_heads / num_kv_heads;
    let scale = 1.0 / (head_dim as f32).sqrt();

    let mut output = vec![0.0f32; seq_len * num_heads * head_dim];

    for h in 0..num_heads {
        let kv_h = h / heads_per_kv; // GQA: multiple Q heads share one KV head

        for s in 0..seq_len {
            // Compute attention scores: Q[s,h] @ K[:kv_len,kv_h]^T
            let mut scores: Vec<f32> = (0..kv_len)
                .map(|t| {
                    if causal && t > (start_pos + s) {
                        return f32::NEG_INFINITY;
                    }
                    let mut dot = 0.0f32;
                    for d in 0..head_dim {
                        let q_val = q.data[(s * num_heads + h) * head_dim + d];
                        let k_val = k.data[(t * num_kv_heads + kv_h) * head_dim + d];
                        dot += q_val * k_val;
                    }
                    dot * scale
                })
                .collect();

            // Softmax
            let max_score = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
            let mut exp_sum = 0.0f32;
            for score in &mut scores {
                *score = (*score - max_score).exp();
                exp_sum += *score;
            }
            if exp_sum > 0.0 {
                for score in &mut scores {
                    *score /= exp_sum;
                }
            }

            // Weighted sum of V
            for d in 0..head_dim {
                let weighted: f32 = scores
                    .iter()
                    .enumerate()
                    .map(|(t, &score)| score * v.data[(t * num_kv_heads + kv_h) * head_dim + d])
                    .sum();
                output[(s * num_heads + h) * head_dim + d] = weighted;
            }
        }
    }

    Tensor {
        data: output,
        shape: vec![seq_len, num_heads * head_dim],
    }
}

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

    #[test]
    fn test_attention_output_shape() {
        let seq = 3;
        let heads = 4;
        let kv_heads = 2;
        let dim = 8;
        let q = Tensor::from_data(vec![0.1; seq * heads * dim], vec![seq, heads * dim]).unwrap();
        let k = Tensor::from_data(vec![0.1; seq * kv_heads * dim], vec![seq, kv_heads * dim]).unwrap();
        let v = Tensor::from_data(vec![0.1; seq * kv_heads * dim], vec![seq, kv_heads * dim]).unwrap();

        let out = attention(&q, &k, &v, &AttentionParams {
            num_heads: heads, num_kv_heads: kv_heads, head_dim: dim,
            seq_len: seq, kv_len: seq, causal: true, start_pos: 0,
        });
        assert_eq!(out.shape, vec![seq, heads * dim]);
    }

    #[test]
    fn test_attention_causal_mask() {
        let heads = 1;
        let dim = 2;
        let seq = 2;
        let q = Tensor::from_data(vec![1.0, 0.0, 0.0, 1.0], vec![2, 2]).unwrap();
        let k = Tensor::from_data(vec![1.0, 0.0, 0.0, 1.0], vec![2, 2]).unwrap();
        let v = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();

        let out = attention(&q, &k, &v, &AttentionParams {
            num_heads: heads, num_kv_heads: heads, head_dim: dim,
            seq_len: seq, kv_len: seq, causal: true, start_pos: 0,
        });
        // Token 0 with causal mask can only see token 0 → output = V[0]
        assert!((out.data[0] - 1.0).abs() < 1e-4);
        assert!((out.data[1] - 2.0).abs() < 1e-4);
    }

    #[test]
    fn test_gqa_heads_share_kv() {
        let dim = 2;
        let q = Tensor::from_data(vec![1.0; 4 * dim], vec![1, 4 * dim]).unwrap();
        let k = Tensor::from_data(vec![1.0; 2 * dim], vec![1, 2 * dim]).unwrap();
        let v = Tensor::from_data(
            vec![1.0, 2.0, 3.0, 4.0],
            vec![1, 2 * dim],
        ).unwrap();

        let out = attention(&q, &k, &v, &AttentionParams {
            num_heads: 4, num_kv_heads: 2, head_dim: dim,
            seq_len: 1, kv_len: 1, causal: false, start_pos: 0,
        });
        // Heads 0,1 share KV head 0 → V[0]=[1,2]
        assert!((out.data[0] - 1.0).abs() < 1e-4);
        assert!((out.data[1] - 2.0).abs() < 1e-4);
        assert!((out.data[2] - 1.0).abs() < 1e-4);
        assert!((out.data[3] - 2.0).abs() < 1e-4);
        // Heads 2,3 share KV head 1 → V[1]=[3,4]
        assert!((out.data[4] - 3.0).abs() < 1e-4);
        assert!((out.data[5] - 4.0).abs() < 1e-4);
    }
}
