// stdlib/crypto/sample.sy — Sampling functions for Kyber-768
//
// CBD: Centered Binomial Distribution (eta=2) for noise polynomials
// Parse: Rejection sampling from SHAKE-128 for matrix A
// PRF: SHAKE-256 based pseudo-random function for deterministic noise
//
// References:
//   - FIPS 203 Algorithm 2 (SampleNTT) and Algorithm 3 (SamplePolyCBD)
//   - Kyber reference: cbd.c, indcpa.c
//
// Aislinn McCubbins, Apr 13 2026

use crypto::fieldq

// ── CBD_eta2: Centered Binomial Distribution with eta=2 ──────────
// Input: 128-byte buffer (Tensor[128] of f32 values 0-255)
// Output: 256-coefficient polynomial (Tensor[256])
// Each coefficient is in {-2, -1, 0, 1, 2}
//
// For each coefficient i:
//   Take 4 consecutive bits from the buffer
//   a = bit0 + bit1 (sum of first 2 bits)
//   b = bit2 + bit3 (sum of next 2 bits)
//   coeff[i] = a - b

fn cbd_eta2(buf: Tensor) -> Tensor {
    let mut poly = t_zeros(256)
    let mut i = 0
    while i < 256 {
        // Each coefficient uses 4 bits = half a byte
        // Coefficient i uses bits [4*i .. 4*i+3]
        // Byte index = (4*i) / 8 = i/2
        // Bit offset within byte = (4*i) % 8 = 0 or 4
        let byte_idx = i / 2
        let byte_val = t_index_get(buf, byte_idx)
        let mut bits = 0
        if i % 2 == 0 {
            // Low 4 bits of the byte
            bits = bit_and(byte_val, 15)
        } else {
            // High 4 bits of the byte
            bits = shr(bit_and(byte_val, 240), 4)
        }
        // a = sum of bits 0,1; b = sum of bits 2,3
        let b0 = bit_and(bits, 1)
        let b1 = bit_and(shr(bits, 1), 1)
        let b2 = bit_and(shr(bits, 2), 1)
        let b3 = bit_and(shr(bits, 3), 1)
        let a = b0 + b1
        let b = b2 + b3
        let coeff = a - b
        poly = t_index_set(poly, i, coeff + 0.0)
        i = i + 1
    }
    return poly
}

// ── PRF: SHAKE-256 based pseudo-random function ──────────────────
// PRF(seed, nonce) = SHAKE-256(seed || nonce_byte, output_len)
// Used for deterministic noise generation in keygen and encryption.

fn kyber_prf(seed: Tensor, nonce: Int, out_len: Int) -> Tensor {
    let nonce_byte = t_zeros(1)
    let nonce_t = t_index_set(nonce_byte, 0, nonce + 0.0)
    let input = t_cat(seed, nonce_t)
    return shake_256(input, out_len)
}

// ── Sample noise polynomial via CBD ──────────────────────────────
// Combines PRF + CBD: generates a noise polynomial from seed + nonce.
// For Kyber-768 with eta=2: needs 2*2*256/8 = 128 bytes from PRF.

fn sample_noise_cbd2(seed: Tensor, nonce: Int) -> Tensor {
    let buf = kyber_prf(seed, nonce, 128)
    return cbd_eta2(buf)
}

// ── Sample matrix A element via rejection sampling ───────────────
// XOF(rho, i, j) = SHAKE-128(rho || i_byte || j_byte)
// Parse 3 bytes at a time into two 12-bit candidates.
// Accept if < q (3329). Reject otherwise. Repeat until 256 coefficients.
//
// FIPS 203 Algorithm 2 (SampleNTT) — output is already in NTT domain.

fn sample_ntt(rho: Tensor, i: Int, j: Int) -> Tensor {
    // Build XOF input: rho (32 bytes) || i || j
    let mut seed = t_zeros(34)
    let mut si = 0
    while si < 32 {
        seed = t_index_set(seed, si, t_index_get(rho, si))
        si = si + 1
    }
    seed = t_index_set(seed, 32, i + 0.0)
    seed = t_index_set(seed, 33, j + 0.0)

    // Request enough bytes for rejection sampling.
    // Expected ~384 bytes needed; request 672 (4 SHAKE blocks) for safety.
    let xof_out = shake_128(seed, 672)

    let mut poly = t_zeros(256)
    let mut coeff_idx = 0
    let mut byte_idx = 0

    while coeff_idx < 256 {
        if byte_idx + 2 >= 672 {
            // Ran out of bytes (extremely unlikely with 672)
            // Fill remaining with 0 as a fallback
            coeff_idx = 256
        } else {
            let b0 = t_index_get(xof_out, byte_idx)
            let b1 = t_index_get(xof_out, byte_idx + 1)
            let b2 = t_index_get(xof_out, byte_idx + 2)
            byte_idx = byte_idx + 3

            // Parse two 12-bit values from 3 bytes
            // d1 = b0 + 256 * (b1 & 0x0F)
            // d2 = (b1 >> 4) + 16 * b2
            let d1 = bit_and(b0, 255) + 256 * bit_and(b1, 15)
            let d2 = shr(bit_and(b1, 255), 4) + 16 * bit_and(b2, 255)

            if d1 < fq_q() {
                if coeff_idx < 256 {
                    poly = t_index_set(poly, coeff_idx, d1 + 0.0)
                    coeff_idx = coeff_idx + 1
                }
            }
            if d2 < fq_q() {
                if coeff_idx < 256 {
                    poly = t_index_set(poly, coeff_idx, d2 + 0.0)
                    coeff_idx = coeff_idx + 1
                }
            }
        }
    }
    return poly
}

// ── Self-test: CBD distribution check ────────────────────────────
fn sample_self_test() -> Int {
    // Generate a noise polynomial and verify all coefficients are in [-2, 2]
    let seed = secure_random_bytes(32)
    let noise = sample_noise_cbd2(seed, 0)

    let mut i = 0
    let mut errors = 0
    let mut min_val = 100
    let mut max_val = -100
    while i < 256 {
        let v = t_index_get(noise, i)
        if v < -2.0 || v > 2.0 {
            errors = errors + 1
        }
        if v < min_val { min_val = v }
        if v > max_val { max_val = v }
        i = i + 1
    }

    // Check that we get SOME variation (not all zeros)
    if min_val == max_val {
        return 2
    }

    return errors
}
