//! INT8 Matrix Multiplication with i32 Accumulation
//!
//! Provides a reusable CPU reference implementation of INT8 GEMM for use by
//! `SmoothQuantizedLayer::forward_raw` and as a future wiring point for GPU kernels.
//!
//! # INT8 Range
//!
//! This module uses the **symmetric** INT8 range **−127…127**, matching SmoothQuant's
//! convention.  The full INT8 range −128…127 is intentionally avoided so that the
//! absolute maximum value is the same on both sides of zero (127), which simplifies
//! per-tensor scale computation as `scale = max_abs / 127.0`.
//!
//! # Accumulation
//!
//! All inner products accumulate into `i32` to avoid overflow.  The maximum
//! theoretical dot-product magnitude for vectors of length N with all elements
//! at ±127 is `127 * 127 * N = 16129 * N`.  An `i32` saturates at ~2.1 × 10⁹,
//! so it handles `N ≤ 131_072` safely — well beyond any realistic hidden dimension.
//!
//! # Dequantisation Formula
//!
//! Given:
//! - `acc_i32`      : INT8 dot product accumulated into i32
//! - `act_scale`    : f32 scale derived as `max(|smoothed activations|) / 127.0`
//! - `weight_scale` : f32 scale stored in `QuantizedTensorData::scale`
//! - `zp`          : i8 zero-point stored in `QuantizedTensorData::zero_point`
//!
//! The output element is:
//!
//! ```text
//! output = (sum_k  act_q[k] * (weight_q[k] - zp))  *  act_scale  *  weight_scale
//! ```
//!
//! where the inner sum was accumulated in i32 before being cast to f32.

use crate::{
    Result,
    error::{CompressionError, QuantizationError},
};

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// INT8 matrix-vector product with i32 accumulation.
///
/// Computes `output[out_row] = Σ_k weights[out_row * in_features + k] * activations[k]`
/// for all output rows, accumulating into i32 to avoid overflow.
///
/// # Arguments
///
/// * `weights`     – Row-major INT8 weight matrix, shape `[out_features, in_features]`.
/// * `activations` – INT8 activation vector, length `in_features`.
/// * `out_features` – Number of output rows.
/// * `in_features`  – Number of input columns.
///
/// # Returns
///
/// `Vec<i32>` of length `out_features`.
///
/// # Errors
///
/// Returns an error if the slice lengths are inconsistent with the declared dimensions.
///
/// # Example
///
/// ```rust
/// use rtx_compress::quantization::int8_matmul::int8_matvec;
///
/// // Identity-like: W = [[1,0],[0,1]], act = [3,5]
/// let weights = vec![1i8, 0, 0, 1];
/// let act     = vec![3i8, 5];
/// let out     = int8_matvec(&weights, &act, 2, 2).unwrap();
/// assert_eq!(out, vec![3i32, 5i32]);
/// ```
pub fn int8_matvec(
    weights: &[i8],
    activations: &[i8],
    out_features: usize,
    in_features: usize,
) -> Result<Vec<i32>> {
    if weights.len() != out_features * in_features {
        return Err(CompressionError::Quantization(
            QuantizationError::InvalidConfig(format!(
                "int8_matvec: weight slice length {} does not match \
                 out_features={} * in_features={}",
                weights.len(),
                out_features,
                in_features
            )),
        ));
    }
    if activations.len() != in_features {
        return Err(CompressionError::Quantization(
            QuantizationError::InvalidConfig(format!(
                "int8_matvec: activation length {} does not match in_features={}",
                activations.len(),
                in_features
            )),
        ));
    }

    let mut output = vec![0i32; out_features];
    for out_row in 0..out_features {
        let row_offset = out_row * in_features;
        let mut acc = 0i32;
        for k in 0..in_features {
            acc += (weights[row_offset + k] as i32) * (activations[k] as i32);
        }
        output[out_row] = acc;
    }
    Ok(output)
}

/// Batched INT8 GEMM: `weights [out × in]` × `activations [batch × in]` → `[batch × out]`.
///
/// Output is row-major `[batch_size, out_features]` with i32 elements.
///
/// # Arguments
///
/// * `weights`      – Row-major INT8 weight matrix, shape `[out_features, in_features]`.
/// * `activations`  – Row-major INT8 activation matrix, shape `[batch_size, in_features]`.
/// * `batch_size`   – Number of input rows.
/// * `out_features` – Number of output rows.
/// * `in_features`  – Number of input columns (must match both slices).
///
/// # Returns
///
/// Row-major `Vec<i32>` of shape `[batch_size, out_features]`.
///
/// # Errors
///
/// Returns an error if slice lengths are inconsistent with declared dimensions.
///
/// # Example
///
/// ```rust
/// use rtx_compress::quantization::int8_matmul::int8_gemm;
///
/// // W = [[1,2],[3,4]], acts = [[1,0],[0,1]] (batch=2)
/// let weights = vec![1i8, 2, 3, 4];
/// let acts    = vec![1i8, 0, 0, 1];
/// let out     = int8_gemm(&weights, &acts, 2, 2, 2).unwrap();
/// // batch 0: [1*1+2*0, 3*1+4*0] = [1, 3]
/// // batch 1: [1*0+2*1, 3*0+4*1] = [2, 4]
/// assert_eq!(out, vec![1i32, 3, 2, 4]);
/// ```
pub fn int8_gemm(
    weights: &[i8],
    activations: &[i8],
    batch_size: usize,
    out_features: usize,
    in_features: usize,
) -> Result<Vec<i32>> {
    if weights.len() != out_features * in_features {
        return Err(CompressionError::Quantization(
            QuantizationError::InvalidConfig(format!(
                "int8_gemm: weight slice length {} does not match \
                 out_features={} * in_features={}",
                weights.len(),
                out_features,
                in_features
            )),
        ));
    }
    if activations.len() != batch_size * in_features {
        return Err(CompressionError::Quantization(
            QuantizationError::InvalidConfig(format!(
                "int8_gemm: activation slice length {} does not match \
                 batch_size={} * in_features={}",
                activations.len(),
                batch_size,
                in_features
            )),
        ));
    }

    let mut output = vec![0i32; batch_size * out_features];
    for b in 0..batch_size {
        let act_offset = b * in_features;
        for out_row in 0..out_features {
            let weight_offset = out_row * in_features;
            let mut acc = 0i32;
            for k in 0..in_features {
                acc += (weights[weight_offset + k] as i32) * (activations[act_offset + k] as i32);
            }
            output[b * out_features + out_row] = acc;
        }
    }
    Ok(output)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    // -----------------------------------------------------------------------
    // int8_matvec tests
    // -----------------------------------------------------------------------

    /// Weight matrix = identity → output equals activation vector.
    #[test]
    fn test_int8_matvec_identity_weights() {
        // W = [[1,0,0],[0,1,0],[0,0,1]] (3×3 identity)
        let weights: Vec<i8> = vec![1, 0, 0, 0, 1, 0, 0, 0, 1];
        let activations: Vec<i8> = vec![7, -3, 5];
        let out = int8_matvec(&weights, &activations, 3, 3).unwrap();
        assert_eq!(out, vec![7i32, -3, 5]);
    }

    /// All-zero activations produce a zero output regardless of weights.
    #[test]
    fn test_int8_matvec_zero_activation() {
        let weights: Vec<i8> = vec![1, 2, 3, 4, 5, 6];
        let activations: Vec<i8> = vec![0, 0, 0];
        let out = int8_matvec(&weights, &activations, 2, 3).unwrap();
        assert_eq!(out, vec![0i32, 0]);
    }

    /// Shape mismatch between declared in_features and activation slice returns Err.
    #[test]
    fn test_int8_matvec_shape_mismatch() {
        // Weight matrix is 2×4 (8 elements), but activations has 3 elements, not 4.
        let weights: Vec<i8> = vec![1, 2, 3, 4, 5, 6, 7, 8];
        let activations: Vec<i8> = vec![1, 2, 3]; // wrong length
        let result = int8_matvec(&weights, &activations, 2, 4);
        assert!(
            result.is_err(),
            "expected Err on in_features mismatch, got Ok"
        );
    }

    // -----------------------------------------------------------------------
    // int8_gemm tests
    // -----------------------------------------------------------------------

    /// batch_size=2 with distinct activation rows produces independent outputs.
    #[test]
    fn test_int8_gemm_batch_size_2() {
        // W = [[1,1],[1,1]] (all-ones 2×2), batch of 2: act0=[2,3], act1=[10,-1]
        let weights: Vec<i8> = vec![1, 1, 1, 1];
        // Row-major [batch=2, in=2]: [act0_col0, act0_col1, act1_col0, act1_col1]
        let activations: Vec<i8> = vec![2, 3, 10, -1];
        let out = int8_gemm(&weights, &activations, 2, 2, 2).unwrap();
        // batch 0, out_row 0: 1*2 + 1*3 = 5
        // batch 0, out_row 1: 1*2 + 1*3 = 5
        // batch 1, out_row 0: 1*10 + 1*(-1) = 9
        // batch 1, out_row 1: 1*10 + 1*(-1) = 9
        assert_eq!(out, vec![5i32, 5, 9, 9]);
    }

    /// Known-value accumulation check: W=[[1,2],[3,4]], acts=[[1,0],[0,1]].
    ///
    /// Output layout is [batch, out_features]:
    ///   batch 0: [1*1+2*0, 3*1+4*0] = [1, 3]
    ///   batch 1: [1*0+2*1, 3*0+4*1] = [2, 4]
    #[test]
    fn test_int8_gemm_accumulation_correct() {
        let weights: Vec<i8> = vec![1, 2, 3, 4]; // [out=2, in=2]
        let activations: Vec<i8> = vec![1, 0, 0, 1]; // [batch=2, in=2]
        let out = int8_gemm(&weights, &activations, 2, 2, 2).unwrap();
        assert_eq!(out, vec![1i32, 3, 2, 4]);
    }

    /// Max INT8 values (±127) accumulate into i32 without overflow.
    ///
    /// For in_features=256, the maximum dot-product is 127 * 127 * 256 = 4,128,256,
    /// well within i32::MAX (2,147,483,647).
    #[test]
    fn test_int8_overflow_clamped() {
        let n = 256usize;
        // Both weights and activations all at +127.
        let weights: Vec<i8> = vec![127i8; n];
        let activations: Vec<i8> = vec![127i8; n];
        let out = int8_matvec(&weights, &activations, 1, n).unwrap();
        let expected = 127i32 * 127 * n as i32; // = 4_128_256
        assert_eq!(out[0], expected);
        // Confirm no silent truncation to i8.
        assert!(out[0] > i8::MAX as i32, "accumulation stayed in i32");
    }

    /// in_features mismatch between weight and activation slice lengths returns Err.
    #[test]
    fn test_int8_matvec_weight_shape_mismatch() {
        // Declare out=2, in=4 but supply only 6 weight elements (≠ 8).
        let weights: Vec<i8> = vec![1, 2, 3, 4, 5, 6];
        let activations: Vec<i8> = vec![1, 2, 3, 4];
        let result = int8_matvec(&weights, &activations, 2, 4);
        assert!(
            result.is_err(),
            "expected Err on weight slice length mismatch"
        );
    }
}