// stdlib/crypto/fieldq.sy — arithmetic in Z_q for Kyber-768 (q = 3329)
//
// This is the foundation layer for Phase 3. Every polynomial operation
// Kyber performs — the NTT, pointwise multiplication, compression, noise
// sampling — is a sequence of mod-q operations on 16-bit coefficients.
// All non-trivial operations go through either Montgomery or Barrett
// reduction because naive `% q` is too slow in hot loops.
//
// Today this is pure Scyla with Int arithmetic, fine for correctness
// testing. When the IntTensor variant lands (§1.4 of CARDEA_PHASE3_PREREQS),
// we'll add vectorized versions operating on whole polynomials.
//
// References:
//   - FIPS 203 ML-KEM (NIST, Aug 2024)
//   - ml-kem crate 0.2.3 (RustCrypto)
//   - Kyber reference C implementation
//
// Aislinn McCubbins, Apr 11 2026

// Bitwise ops (bit_and, bit_or, bit_xor, shl, shr) are now VM builtins.
// No import needed — they handle negative numbers correctly unlike the
// arithmetic fallbacks in bitops.sy.

// ── Constants ────────────────────────────────────────────
// These are public so other modules (ntt.sy, cbd.sy, poly.sy) can see them.

fn fq_q() -> Int { return 3329 }

// 2^16 mod q. Used to enter Montgomery form: to_mont(a) = mont_reduce(a * MONT_R2)
fn fq_mont() -> Int { return 2285 }

// (2^16)^2 mod q = 2285 * 2285 mod 3329. Used for to_mont.
// Verification: 2285 * 2285 = 5221225; 5221225 % 3329 = 1353.
fn fq_mont_r2() -> Int { return 1353 }

// q^(-1) mod 2^16. Verified: 3329 * 62209 mod 65536 == 1.
fn fq_qinv() -> Int { return 62209 }

// Primitive 256th root of unity mod q. The NTT twiddle factor base.
fn fq_zeta_base() -> Int { return 17 }

// Polynomial length
fn fq_n() -> Int { return 256 }

// ── Basic arithmetic (no reduction) ──────────────────────
// These assume inputs are already in canonical range [0, q) and return
// results in [0, q). Cost: one conditional subtract. Good for add/sub
// in hot NTT loops.

fn fq_add(a: Int, b: Int) -> Int {
    let mut r = a + b
    if r >= fq_q() {
        r = r - fq_q()
    }
    return r
}

fn fq_sub(a: Int, b: Int) -> Int {
    let mut r = a - b
    if r < 0 {
        r = r + fq_q()
    }
    return r
}

// ── Barrett reduction ────────────────────────────────────
// Reduces a 26-bit non-negative value into [0, q).
// Formula from Kyber reference:
//     v = floor((2^26 + q/2) / q) = 20159
//     t = (v * a + 2^25) >> 26
//     r = a - t * q
//     if r >= q then r -= q
// We substitute >>26 with /(2^26) and >>25 with /(2^25).

fn fq_barrett_v() -> Int { return 20159 }

fn fq_barrett(a: Int) -> Int {
    let t = shr(fq_barrett_v() * a + 33554432, 26)
    let mut r = a - t * fq_q()
    if r >= fq_q() {
        r = r - fq_q()
    }
    if r < 0 {
        r = r + fq_q()
    }
    return r
}

// ── Montgomery reduction ─────────────────────────────────
// Given a 32-bit signed input, compute (a * R^(-1)) mod q where R = 2^16.
// Output range is [-q, q) (NOT canonical). Caller normalizes if needed.
//
// Reference C:
//   int16_t t = (int16_t)(a * QINV);
//   t = (a - (int32_t)t * Q) >> 16;
//   return t;
//
// We simulate (int16_t) truncation: take `(a * QINV) mod 2^16` and if the
// result is >= 2^15, subtract 2^16 to get a signed 16-bit value.

fn fq_mont_reduce(a: Int) -> Int {
    // Low 16 bits of a * QINV, as signed int16.
    let mut t = bit_and(a * fq_qinv(), 65535)
    if t >= 32768 {
        t = t - 65536
    }

    // (a - t * q) >> 16. The low 16 bits of (a - t*q) are zero by
    // construction, so arithmetic right shift is exact.
    return shr(a - t * fq_q(), 16)
}

// Multiply two values with Montgomery reduction. If inputs are in
// Montgomery form (a*R, b*R), output is a*b*R.
fn fq_mul(a: Int, b: Int) -> Int {
    return fq_mont_reduce(a * b)
}

// Normalize to canonical range [0, q)
fn fq_normalize(a: Int) -> Int {
    let mut r = a
    if r < 0 { r = r + fq_q() }
    if r >= fq_q() { r = r - fq_q() }
    return r
}

// Enter Montgomery form: a -> a * R mod q
fn fq_to_mont(a: Int) -> Int {
    return fq_mont_reduce(a * fq_mont_r2())
}

// Leave Montgomery form: a*R -> a, normalized to [0, q)
fn fq_from_mont(a: Int) -> Int {
    return fq_normalize(fq_mont_reduce(a))
}

// ── Power (square-and-multiply) ──────────────────────────
// Used to build the ZETAS table: zetas[i] = zeta_base^bitrev7(i) mod q
// (in Montgomery form). Only called at init, so O(log exp) is fine.

fn fq_pow(base: Int, exp: Int) -> Int {
    let mut result = 1
    let mut b = base
    let mut e = exp
    while e > 0 {
        if e % 2 == 1 {
            result = (result * b) % fq_q()
        }
        b = (b * b) % fq_q()
        e = e / 2
    }
    return result
}

// ── Self-test helpers ────────────────────────────────────
// Sanity checks that Phase 3 test harness can call to verify the fallback
// implementation before swapping in VM builtins. These are NOT crypto
// tests — just arithmetic invariants.

fn fq_self_test() -> Int {
    // 1. mont round-trip: from_mont(to_mont(x)) == x for x in [0, q)
    let x = 1234
    let round = fq_from_mont(fq_to_mont(x))
    if round != x {
        return 1
    }

    // 2. add/sub are inverses
    let s = fq_add(100, 200)
    if fq_sub(s, 200) != 100 {
        return 2
    }

    // 3. mul identity in Montgomery form: a * R * 1 * R^-1 = a
    // to_mont(a) * 1_mont / R = a_mont, then from_mont gives a
    let a_m = fq_to_mont(42)
    let one_m = fq_to_mont(1)
    let prod = fq_mul(a_m, one_m)
    if fq_from_mont(prod) != 42 {
        return 3
    }

    // 4. barrett matches % for small values
    if fq_barrett(5000) != (5000 % fq_q()) {
        return 4
    }

    // 5. pow: zeta_base^256 == 1 mod q (256th root of unity)
    if fq_pow(fq_zeta_base(), 256) != 1 {
        return 5
    }

    // 6. pow: zeta_base^128 == q - 1 (primitive root, not a lower-order one)
    if fq_pow(fq_zeta_base(), 128) != fq_q() - 1 {
        return 6
    }

    return 0
}
