//! Grouped Query Attention (GQA) KV head expansion for inference.
//!
//! Modern LLMs such as Llama 2/3, Mistral, Gemma, and Qwen use GQA where
//! `num_kv_heads < num_q_heads`.  During the attention forward pass each KV
//! head is shared by `queries_per_group = num_q_heads / num_kv_heads` query
//! heads, eliminating the need to store a full per-head KV cache.
//!
//! This module provides:
//! - [`GqaConfig`] — validated configuration object.
//! - [`expand_kv_heads`] — repeat/tile KV tensor to match Q head count.
//! - [`kv_head_for_q`] — zero-cost index mapping from Q head → KV head.
//! - [`gqa_attention_cpu`] — single-batch GQA forward pass on the CPU.
//!
//! # Layout convention
//!
//! Tensors are represented as flat `&[f32]` / `Vec<f32>` in
//! **row-major (C-contiguous)** order.  The logical shape is written in
//! square brackets throughout the docs.
//!
//! ```text
//! KV tensor: [batch, num_kv_heads, seq_len, head_dim]
//! Q  tensor: [num_q_heads, seq_len_q, head_dim]        (single batch item)
//! ```

use thiserror::Error;

// ──────────────────────────────────────────────────────────────────────────────
// Error type
// ──────────────────────────────────────────────────────────────────────────────

/// Errors that can occur when constructing or using a [`GqaConfig`].
#[derive(Debug, Error, PartialEq, Eq)]
pub enum GqaError {
    /// `num_q_heads` is not evenly divisible by `num_kv_heads`.
    #[error("num_q_heads ({q}) must be divisible by num_kv_heads ({kv})")]
    IndivisibleHeads { q: usize, kv: usize },

    /// Either head count is zero.
    #[error("num_q_heads and num_kv_heads must be > 0")]
    ZeroHeads,

    /// `num_kv_heads` is larger than `num_q_heads`, which is not valid.
    #[error("num_kv_heads ({kv}) cannot exceed num_q_heads ({q})")]
    KvExceedsQ { q: usize, kv: usize },

    /// `head_dim` is zero.
    #[error("head_dim must be > 0")]
    ZeroHeadDim,
}

// ──────────────────────────────────────────────────────────────────────────────
// Configuration
// ──────────────────────────────────────────────────────────────────────────────

/// Configuration for Grouped Query Attention.
///
/// # Invariants (upheld by [`GqaConfig::new`])
/// - `num_q_heads > 0`
/// - `num_kv_heads > 0`
/// - `num_kv_heads <= num_q_heads`
/// - `num_q_heads % num_kv_heads == 0`
/// - `head_dim > 0`
///
/// # Examples
///
/// ```
/// use rtx_inference::gqa::GqaConfig;
///
/// // Llama 3 8B: 32 Q heads, 8 KV heads, head_dim 128.
/// let cfg = GqaConfig::new(32, 8, 128).unwrap();
/// assert_eq!(cfg.queries_per_group(), 4);
/// assert!(cfg.is_grouped());
/// assert!(!cfg.is_mqa());
/// assert!((cfg.kv_memory_ratio() - 0.25).abs() < 1e-6);
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GqaConfig {
    /// Total number of query heads.
    pub num_q_heads: usize,
    /// Total number of key/value heads (≤ `num_q_heads`).
    pub num_kv_heads: usize,
    /// Dimension of each head.
    pub head_dim: usize,
}

impl GqaConfig {
    /// Construct a validated [`GqaConfig`].
    ///
    /// # Errors
    ///
    /// Returns [`GqaError`] when any invariant is violated.
    pub fn new(num_q_heads: usize, num_kv_heads: usize, head_dim: usize) -> Result<Self, GqaError> {
        if num_q_heads == 0 || num_kv_heads == 0 {
            return Err(GqaError::ZeroHeads);
        }
        if head_dim == 0 {
            return Err(GqaError::ZeroHeadDim);
        }
        if num_kv_heads > num_q_heads {
            return Err(GqaError::KvExceedsQ {
                q: num_q_heads,
                kv: num_kv_heads,
            });
        }
        if num_q_heads % num_kv_heads != 0 {
            return Err(GqaError::IndivisibleHeads {
                q: num_q_heads,
                kv: num_kv_heads,
            });
        }
        Ok(Self {
            num_q_heads,
            num_kv_heads,
            head_dim,
        })
    }

    /// Number of query heads that share one KV head.
    ///
    /// Equals `num_q_heads / num_kv_heads`.  For standard MHA this is `1`.
    #[inline]
    pub fn queries_per_group(&self) -> usize {
        self.num_q_heads / self.num_kv_heads
    }

    /// Returns `true` when GQA is active (`num_kv_heads < num_q_heads`).
    #[inline]
    pub fn is_grouped(&self) -> bool {
        self.num_kv_heads < self.num_q_heads
    }

    /// Returns `true` for Multi-Query Attention (single KV head for all Q heads).
    #[inline]
    pub fn is_mqa(&self) -> bool {
        self.num_kv_heads == 1
    }

    /// Ratio of KV memory relative to full MHA.
    ///
    /// E.g. 8 KV heads / 32 Q heads = 0.25  →  75 % memory saving.
    #[inline]
    pub fn kv_memory_ratio(&self) -> f32 {
        self.num_kv_heads as f32 / self.num_q_heads as f32
    }
}

// ──────────────────────────────────────────────────────────────────────────────
// Index mapping
// ──────────────────────────────────────────────────────────────────────────────

/// Map a query head index to its corresponding KV head index.
///
/// The mapping is: `kv_head = q_head / queries_per_group`.
///
/// This is a zero-cost inline function suitable for tight inner loops.
///
/// # Examples
///
/// ```
/// use rtx_inference::gqa::kv_head_for_q;
///
/// // 8 Q heads, 2 KV heads → queries_per_group = 4
/// assert_eq!(kv_head_for_q(0, 4), 0);
/// assert_eq!(kv_head_for_q(3, 4), 0);
/// assert_eq!(kv_head_for_q(4, 4), 1);
/// assert_eq!(kv_head_for_q(7, 4), 1);
/// ```
#[inline]
pub fn kv_head_for_q(q_head: usize, queries_per_group: usize) -> usize {
    q_head / queries_per_group
}

// ──────────────────────────────────────────────────────────────────────────────
// KV head expansion
// ──────────────────────────────────────────────────────────────────────────────

/// Expand KV heads to match the query head count by repetition (tiling).
///
/// Input layout:  `[batch, num_kv_heads, seq_len, head_dim]`
/// Output layout: `[batch, num_q_heads,  seq_len, head_dim]`
///
/// Each KV head `i` is repeated [`GqaConfig::queries_per_group`] times
/// consecutively so that the output can be used directly with a full-MHA
/// kernel.
///
/// # Memory
///
/// Allocates a new `Vec<f32>` of length
/// `batch_size * num_q_heads * seq_len * head_dim`.
///
/// # Panics
///
/// Panics in debug mode if `kv.len()` does not match the expected size
/// `batch_size * config.num_kv_heads * seq_len * config.head_dim`.
///
/// # Examples
///
/// ```
/// use rtx_inference::gqa::{GqaConfig, expand_kv_heads};
///
/// let cfg = GqaConfig::new(8, 2, 4).unwrap();
/// // KV: [1, 2, 3, 4]  (batch=1, num_kv=2, seq=3, head_dim=4)
/// let kv: Vec<f32> = (0..24).map(|x| x as f32).collect();
/// let out = expand_kv_heads(&kv, &cfg, 1, 3);
/// assert_eq!(out.len(), 1 * 8 * 3 * 4);
/// ```
pub fn expand_kv_heads(
    kv: &[f32],
    config: &GqaConfig,
    batch_size: usize,
    seq_len: usize,
) -> Vec<f32> {
    let kv_head_stride = seq_len * config.head_dim;
    let kv_batch_stride = config.num_kv_heads * kv_head_stride;
    let q_head_stride = seq_len * config.head_dim;
    let q_batch_stride = config.num_q_heads * q_head_stride;
    let gpg = config.queries_per_group();

    debug_assert_eq!(
        kv.len(),
        batch_size * kv_batch_stride,
        "KV slice length mismatch: expected {} got {}",
        batch_size * kv_batch_stride,
        kv.len()
    );

    let mut out = vec![0.0f32; batch_size * q_batch_stride];

    for b in 0..batch_size {
        for q_head in 0..config.num_q_heads {
            let kv_head = kv_head_for_q(q_head, gpg);
            let kv_offset = b * kv_batch_stride + kv_head * kv_head_stride;
            let out_offset = b * q_batch_stride + q_head * q_head_stride;
            // Copy the entire [seq_len, head_dim] slice.
            out[out_offset..out_offset + kv_head_stride]
                .copy_from_slice(&kv[kv_offset..kv_offset + kv_head_stride]);
        }
    }

    out
}

// ──────────────────────────────────────────────────────────────────────────────
// CPU attention forward pass
// ──────────────────────────────────────────────────────────────────────────────

/// Compute GQA scaled dot-product attention for a single batch item on the CPU.
///
/// Tensors are flat row-major slices:
/// - `q`: `[num_q_heads, seq_len_q, head_dim]`
/// - `k`: `[num_kv_heads, seq_len_kv, head_dim]`  *(NOT pre-expanded)*
/// - `v`: `[num_kv_heads, seq_len_kv, head_dim]`  *(NOT pre-expanded)*
///
/// Returns `[num_q_heads, seq_len_q, head_dim]`.
///
/// For each query head `h_q`, the corresponding KV head is
/// `h_q / queries_per_group`.  Attention is computed as:
///
/// ```text
/// scores[h_q, i, j] = dot(Q[h_q, i, :], K[h_kv, j, :]) * scale
/// attn[h_q, i, j]   = softmax(scores[h_q, i, :])[j]   (causal mask applied before softmax)
/// out[h_q, i, :]     = sum_j attn[h_q, i, j] * V[h_kv, j, :]
/// ```
///
/// When `causal = true` positions `j > i` are masked to `-∞` before softmax.
///
/// # Panics
///
/// Panics in debug mode on slice length mismatch.
///
/// # Examples
///
/// ```
/// use rtx_inference::gqa::{GqaConfig, gqa_attention_cpu};
///
/// let cfg = GqaConfig::new(2, 2, 4).unwrap();
/// let scale = (4f32).sqrt().recip();
/// let q: Vec<f32> = (0..16).map(|x| x as f32 * 0.01).collect();
/// let k = q.clone();
/// let v = q.clone();
/// let out = gqa_attention_cpu(&q, &k, &v, &cfg, 2, 2, scale, false);
/// assert_eq!(out.len(), 2 * 2 * 4);
/// ```
pub fn gqa_attention_cpu(
    q: &[f32],
    k: &[f32],
    v: &[f32],
    config: &GqaConfig,
    seq_len_q: usize,
    seq_len_kv: usize,
    scale: f32,
    causal: bool,
) -> Vec<f32> {
    let head_dim = config.head_dim;
    let gpg = config.queries_per_group();

    debug_assert_eq!(q.len(), config.num_q_heads * seq_len_q * head_dim);
    debug_assert_eq!(k.len(), config.num_kv_heads * seq_len_kv * head_dim);
    debug_assert_eq!(v.len(), config.num_kv_heads * seq_len_kv * head_dim);

    // Strides for input tensors.
    let q_head_stride = seq_len_q * head_dim;
    let kv_head_stride = seq_len_kv * head_dim;

    // Output: [num_q_heads, seq_len_q, head_dim]
    let mut output = vec![0.0f32; config.num_q_heads * seq_len_q * head_dim];

    // Scratch buffer for attention scores of one query position: [seq_len_kv]
    let mut scores = vec![0.0f32; seq_len_kv];

    for h_q in 0..config.num_q_heads {
        let h_kv = kv_head_for_q(h_q, gpg);
        let q_base = h_q * q_head_stride;
        let kv_base = h_kv * kv_head_stride;
        let out_base = h_q * q_head_stride;

        for i in 0..seq_len_q {
            // Compute raw dot-product scores for query position i.
            let q_row = &q[q_base + i * head_dim..q_base + i * head_dim + head_dim];

            let max_j = if causal { i + 1 } else { seq_len_kv };

            // Fill scores; positions beyond causal boundary stay at -∞.
            for j in 0..seq_len_kv {
                if j < max_j {
                    let k_row = &k[kv_base + j * head_dim..kv_base + j * head_dim + head_dim];
                    scores[j] = dot(q_row, k_row) * scale;
                } else {
                    scores[j] = f32::NEG_INFINITY;
                }
            }

            // Numerically stable softmax over scores[0..seq_len_kv].
            let max_score = scores[..seq_len_kv]
                .iter()
                .copied()
                .fold(f32::NEG_INFINITY, f32::max);

            let mut sum_exp = 0.0f32;
            for s in &mut scores[..seq_len_kv] {
                *s = (*s - max_score).exp();
                sum_exp += *s;
            }
            // Guard against the degenerate all-masked case.
            let inv_sum = if sum_exp > 0.0 { sum_exp.recip() } else { 0.0 };
            for s in &mut scores[..seq_len_kv] {
                *s *= inv_sum;
            }

            // Weighted sum over value vectors.
            let out_row = &mut output[out_base + i * head_dim..out_base + i * head_dim + head_dim];
            for j in 0..seq_len_kv {
                let attn = scores[j];
                if attn == 0.0 {
                    continue;
                }
                let v_row = &v[kv_base + j * head_dim..kv_base + j * head_dim + head_dim];
                for d in 0..head_dim {
                    out_row[d] += attn * v_row[d];
                }
            }
        }
    }

    output
}

// ──────────────────────────────────────────────────────────────────────────────
// Internal helpers
// ──────────────────────────────────────────────────────────────────────────────

/// Dot product of two equal-length slices.
#[inline]
fn dot(a: &[f32], b: &[f32]) -> f32 {
    debug_assert_eq!(a.len(), b.len());
    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
}

// ──────────────────────────────────────────────────────────────────────────────
// Tests
// ──────────────────────────────────────────────────────────────────────────────

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

    // ── Config validation tests ───────────────────────────────────────────────

    /// MHA (num_q == num_kv) → queries_per_group == 1.
    #[test]
    fn test_mha_config_queries_per_group_one() {
        let cfg = GqaConfig::new(8, 8, 64).unwrap();
        assert_eq!(cfg.queries_per_group(), 1);
        assert!(!cfg.is_grouped());
    }

    /// num_q not divisible by num_kv → IndivisibleHeads error.
    #[test]
    fn test_gqa_config_indivisible_error() {
        let err = GqaConfig::new(7, 3, 64).unwrap_err();
        assert_eq!(err, GqaError::IndivisibleHeads { q: 7, kv: 3 });
    }

    /// num_kv > num_q → KvExceedsQ error.
    #[test]
    fn test_gqa_config_kv_exceeds_q_error() {
        let err = GqaConfig::new(4, 8, 64).unwrap_err();
        assert_eq!(err, GqaError::KvExceedsQ { q: 4, kv: 8 });
    }

    /// num_q == 0 → ZeroHeads error.
    #[test]
    fn test_gqa_config_zero_heads_error() {
        let err = GqaConfig::new(0, 0, 64).unwrap_err();
        assert_eq!(err, GqaError::ZeroHeads);
    }

    /// num_kv == 0 with non-zero num_q → ZeroHeads error.
    #[test]
    fn test_gqa_config_zero_kv_heads_error() {
        let err = GqaConfig::new(8, 0, 64).unwrap_err();
        assert_eq!(err, GqaError::ZeroHeads);
    }

    /// num_kv < num_q → is_grouped() == true.
    #[test]
    fn test_is_grouped_true() {
        let cfg = GqaConfig::new(32, 8, 128).unwrap();
        assert!(cfg.is_grouped());
    }

    /// num_kv == 1 → is_mqa() == true.
    #[test]
    fn test_is_mqa() {
        let cfg = GqaConfig::new(8, 1, 64).unwrap();
        assert!(cfg.is_mqa());
        assert!(cfg.is_grouped());
    }

    /// 8 kv / 32 q → kv_memory_ratio == 0.25.
    #[test]
    fn test_kv_memory_ratio() {
        let cfg = GqaConfig::new(32, 8, 128).unwrap();
        assert!((cfg.kv_memory_ratio() - 0.25_f32).abs() < 1e-6);
    }

    // ── Index mapping tests ───────────────────────────────────────────────────

    /// q_head=5, queries_per_group=4 → kv_head=1.
    #[test]
    fn test_kv_head_for_q_mapping() {
        assert_eq!(kv_head_for_q(5, 4), 1);
    }

    /// Full group boundary checks for queries_per_group=4.
    #[test]
    fn test_kv_head_for_q_boundary() {
        // 8 Q heads, 2 KV heads, gpg=4
        for q in 0..4usize {
            assert_eq!(kv_head_for_q(q, 4), 0);
        }
        for q in 4..8usize {
            assert_eq!(kv_head_for_q(q, 4), 1);
        }
    }

    // ── expand_kv_heads shape test ────────────────────────────────────────────

    /// Input [1, 2, 4, 8] → output [1, 8, 4, 8] for num_q=8, num_kv=2.
    #[test]
    fn test_expand_kv_heads_shape() {
        let cfg = GqaConfig::new(8, 2, 8).unwrap();
        let batch = 1;
        let seq = 4;
        let kv: Vec<f32> = (0..(batch * 2 * seq * 8)).map(|x| x as f32).collect();
        let out = expand_kv_heads(&kv, &cfg, batch, seq);
        assert_eq!(out.len(), batch * 8 * seq * 8);
    }

    /// KV[0] is repeated 4× (first four Q slots), KV[1] repeats for the rest.
    #[test]
    fn test_expand_kv_heads_values() {
        // cfg: 8 Q heads, 2 KV heads, head_dim=4, batch=1, seq=1
        let cfg = GqaConfig::new(8, 2, 4).unwrap();
        // KV head 0: [1,2,3,4], KV head 1: [5,6,7,8]
        let kv: Vec<f32> = vec![1., 2., 3., 4., 5., 6., 7., 8.];
        let out = expand_kv_heads(&kv, &cfg, 1, 1);
        // Expect first 4 Q heads to carry KV[0] and next 4 to carry KV[1].
        for q in 0..4 {
            let slice = &out[q * 4..(q + 1) * 4];
            assert_eq!(
                slice,
                &[1., 2., 3., 4.],
                "Q head {q} should match KV head 0"
            );
        }
        for q in 4..8 {
            let slice = &out[q * 4..(q + 1) * 4];
            assert_eq!(
                slice,
                &[5., 6., 7., 8.],
                "Q head {q} should match KV head 1"
            );
        }
    }

    // ── gqa_attention_cpu output shape ────────────────────────────────────────

    /// Output shape must be [num_q_heads, seq_len_q, head_dim].
    #[test]
    fn test_gqa_attention_output_shape() {
        let cfg = GqaConfig::new(4, 2, 8).unwrap();
        let seq_q = 3;
        let seq_kv = 5;
        let scale = (8f32).sqrt().recip();
        let q: Vec<f32> = vec![0.1; 4 * seq_q * 8];
        let k: Vec<f32> = vec![0.1; 2 * seq_kv * 8];
        let v: Vec<f32> = vec![0.2; 2 * seq_kv * 8];
        let out = gqa_attention_cpu(&q, &k, &v, &cfg, seq_q, seq_kv, scale, false);
        assert_eq!(out.len(), 4 * seq_q * 8);
    }

    /// When num_kv == num_q, gqa_attention_cpu must match naive MHA within 1e-5.
    #[test]
    fn test_gqa_attention_mha_matches_standard() {
        let n_heads = 2;
        let seq = 3;
        let head_dim = 4;
        let cfg = GqaConfig::new(n_heads, n_heads, head_dim).unwrap();
        let scale = (head_dim as f32).sqrt().recip();

        // Simple ascending values for reproducibility.
        let q: Vec<f32> = (0..(n_heads * seq * head_dim))
            .map(|x| x as f32 * 0.05)
            .collect();
        let k = q.clone();
        let v = q.clone();

        let gqa_out = gqa_attention_cpu(&q, &k, &v, &cfg, seq, seq, scale, false);

        // Compute naive MHA reference (same algorithm, but explicit expansion).
        let ref_out = naive_mha_reference(&q, &k, &v, n_heads, seq, head_dim, scale, false);

        assert_eq!(gqa_out.len(), ref_out.len());
        for (a, b) in gqa_out.iter().zip(ref_out.iter()) {
            assert!((a - b).abs() < 1e-5, "GQA/MHA mismatch: {a} vs {b}");
        }
    }

    /// Causal mask: for seq=4, the output should reflect that future tokens are
    /// invisible.  Specifically, position 0 may only attend to key 0, so the
    /// output for position 0 should equal `value[0]` exactly (after softmax of
    /// a single-element distribution = 1.0).
    #[test]
    fn test_gqa_attention_causal_mask() {
        let n_heads = 1;
        let seq = 4;
        let head_dim = 4;
        let cfg = GqaConfig::new(n_heads, n_heads, head_dim).unwrap();
        let scale = (head_dim as f32).sqrt().recip();

        // Distinct value vectors so we can tell which positions contributed.
        // V row j = all j+1 (1.0, 2.0, 3.0, 4.0 for j=0,1,2,3).
        let q: Vec<f32> = vec![0.1; n_heads * seq * head_dim];
        let k: Vec<f32> = vec![0.0; n_heads * seq * head_dim];
        let mut v: Vec<f32> = vec![0.0; n_heads * seq * head_dim];
        for j in 0..seq {
            for d in 0..head_dim {
                v[j * head_dim + d] = (j + 1) as f32;
            }
        }

        let out = gqa_attention_cpu(&q, &k, &v, &cfg, seq, seq, scale, true);

        // Position 0 (q_i=0): can only attend to k_j=0 → attn weight = 1.0 → out = V[0] = [1,1,1,1]
        let pos0 = &out[0..head_dim];
        for &val in pos0 {
            assert!(
                (val - 1.0).abs() < 1e-5,
                "causal pos0: expected 1.0 got {val}"
            );
        }

        // Position 3 (q_i=3): attends to k_j ∈ {0,1,2,3} uniformly (all keys identical).
        // Expected output = mean of V rows = (1+2+3+4)/4 = 2.5.
        let pos3 = &out[3 * head_dim..4 * head_dim];
        for &val in pos3 {
            assert!(
                (val - 2.5).abs() < 1e-4,
                "causal pos3: expected 2.5 got {val}"
            );
        }
    }

    // ── Helper ────────────────────────────────────────────────────────────────

    /// Naive MHA reference (full expansion then standard attention).
    fn naive_mha_reference(
        q: &[f32],
        k: &[f32],
        v: &[f32],
        n_heads: usize,
        seq: usize,
        head_dim: usize,
        scale: f32,
        causal: bool,
    ) -> Vec<f32> {
        let mut out = vec![0.0f32; n_heads * seq * head_dim];
        let head_stride = seq * head_dim;
        let mut scores = vec![0.0f32; seq];

        for h in 0..n_heads {
            let q_base = h * head_stride;
            let k_base = h * head_stride;
            let v_base = h * head_stride;
            let o_base = h * head_stride;

            for i in 0..seq {
                let q_row = &q[q_base + i * head_dim..q_base + i * head_dim + head_dim];
                let max_j = if causal { i + 1 } else { seq };

                for j in 0..seq {
                    if j < max_j {
                        let k_row = &k[k_base + j * head_dim..k_base + j * head_dim + head_dim];
                        scores[j] = dot(q_row, k_row) * scale;
                    } else {
                        scores[j] = f32::NEG_INFINITY;
                    }
                }

                let max_s = scores[..seq]
                    .iter()
                    .copied()
                    .fold(f32::NEG_INFINITY, f32::max);
                let mut sum_e = 0.0f32;
                for s in &mut scores[..seq] {
                    *s = (*s - max_s).exp();
                    sum_e += *s;
                }
                let inv = if sum_e > 0.0 { sum_e.recip() } else { 0.0 };
                for s in &mut scores[..seq] {
                    *s *= inv;
                }

                let o_row = &mut out[o_base + i * head_dim..o_base + i * head_dim + head_dim];
                for j in 0..seq {
                    let a = scores[j];
                    if a == 0.0 {
                        continue;
                    }
                    let v_row = &v[v_base + j * head_dim..v_base + j * head_dim + head_dim];
                    for d in 0..head_dim {
                        o_row[d] += a * v_row[d];
                    }
                }
            }
        }
        out
    }
}