// stdlib/crypto/ntt.sy — Number Theoretic Transform for Kyber-768
//
// The NTT converts a 256-coefficient polynomial in Z_q[X]/(X^256+1)
// to "NTT domain" where polynomial multiplication becomes pointwise.
// This is what makes Kyber fast: O(n log n) instead of O(n^2).
//
// All arithmetic is in Montgomery form via fieldq.sy.
// Zetas table is precomputed from the Kyber reference implementation.
//
// References:
//   - FIPS 203 ML-KEM §4.3 (NTT and NTT^-1)
//   - Kyber reference C: ntt.c
//   - ζ = 17 is the primitive 256th root of unity mod 3329
//
// Aislinn McCubbins, Apr 13 2026

use crypto::fieldq

// ── Zeta table (Montgomery form, bit-reversed order) ──────
// Generated at runtime: zetas[k] = to_mont(17^{bitrev7(k)}) mod 3329.
// Computed once, cached for NTT use.

fn ntt_bitrev7(n: Int) -> Int {
    let mut result = 0
    let mut val = n
    let mut i = 0
    while i < 7 {
        result = result * 2 + val % 2
        val = val / 2
        i = i + 1
    }
    return result
}

fn ntt_zetas() -> Tensor {
    let mut z = t_zeros(128)
    let mut k = 0
    while k < 128 {
        let br = ntt_bitrev7(k)
        let pow = fq_pow(fq_zeta_base(), br)
        let mont = fq_to_mont(pow)
        z = t_index_set(z, k, mont + 0.0)
        k = k + 1
    }
    return z
}

// ── Montgomery multiply (fqmul) ─────────────────────────────────
// Multiply two values in Montgomery form and reduce.
fn fqmul(a: Int, b: Int) -> Int {
    return fq_mont_reduce(a * b)
}

// ── Forward NTT (Cooley-Tukey butterfly) ─────────────────────────
// In-place NTT on a 256-coefficient polynomial stored as a Tensor.
// After NTT, coefficients are in NTT domain (Montgomery form).
//
// Reference C:
//   k = 1;
//   for(len = 128; len >= 2; len >>= 1)
//     for(start = 0; start < 256; start = j + len)
//       zeta = zetas[k++];
//       for(j = start; j < start + len; j++)
//         t = fqmul(zeta, r[j+len]);
//         r[j+len] = r[j] - t;
//         r[j] = r[j] + t;

fn ntt_forward(poly: Tensor) -> Tensor {
    let zetas = ntt_zetas()
    let mut r = poly
    let mut k = 1
    let mut len = 128

    while len >= 2 {
        let mut start = 0
        while start < 256 {
            let zeta = t_index_get(zetas, k)
            k = k + 1
            let mut j = start
            while j < start + len {
                let rj = t_index_get(r, j)
                let rjl = t_index_get(r, j + len)
                let t = fqmul(zeta, rjl)
                r = t_index_set(r, j + len, rj - t)
                r = t_index_set(r, j, rj + t)
                j = j + 1
            }
            start = start + len + len
        }
        len = len / 2
    }
    return r
}

// ── Inverse NTT (Gentleman-Sande butterfly) ──────────────────────
// Converts from NTT domain back to coefficient domain.
// Includes the final multiplication by f = 1441 = mont^2/128.
//
// Reference C:
//   k = 127;
//   for(len = 2; len <= 128; len <<= 1)
//     for(start = 0; start < 256; start = j + len)
//       zeta = zetas[k--];
//       for(j = start; j < start + len; j++)
//         t = r[j];
//         r[j] = barrett_reduce(t + r[j+len]);
//         r[j+len] = r[j+len] - t;
//         r[j+len] = fqmul(zeta, r[j+len]);
//   for(j = 0; j < 256; j++) r[j] = fqmul(r[j], f);

fn ntt_inverse(poly: Tensor) -> Tensor {
    let zetas = ntt_zetas()
    let mut r = poly
    let mut k = 127
    let mut len = 2

    while len <= 128 {
        let mut start = 0
        while start < 256 {
            let zeta = t_index_get(zetas, k)
            k = k - 1
            let mut j = start
            while j < start + len {
                let t = t_index_get(r, j)
                let rjl = t_index_get(r, j + len)
                r = t_index_set(r, j, fq_barrett(t + rjl))
                let diff = rjl - t
                r = t_index_set(r, j + len, fqmul(zeta, diff))
                j = j + 1
            }
            start = start + len + len
        }
        len = len * 2
    }

    // Final scaling by f = 1441 (Montgomery form of 128^{-1} mod q)
    let f = 1441
    let mut j = 0
    while j < 256 {
        let v = t_index_get(r, j)
        r = t_index_set(r, j, fqmul(v, f))
        j = j + 1
    }
    return r
}

// ── Pointwise multiplication (basemul) ───────────────────────────
// In NTT domain, multiplication of two polynomials is done pairwise
// on 128 pairs of coefficients. Each pair (a0,a1)*(b0,b1) uses a
// specific zeta for the cross term.
//
// basemul_acc((a0,a1), (b0,b1), zeta) =
//   r0 = a0*b0 + a1*b1*zeta
//   r1 = a0*b1 + a1*b0

fn ntt_basemul(a: Tensor, b: Tensor) -> Tensor {
    let zetas = ntt_zetas()
    let mut r = t_zeros(256)
    let mut i = 0

    while i < 64 {
        let z_idx = 64 + i
        let zeta = t_index_get(zetas, z_idx)

        let j = 4 * i

        // First pair
        let a0 = t_index_get(a, j)
        let a1 = t_index_get(a, j + 1)
        let b0 = t_index_get(b, j)
        let b1 = t_index_get(b, j + 1)
        r = t_index_set(r, j, fqmul(a0, b0) + fqmul(fqmul(a1, b1), zeta))
        r = t_index_set(r, j + 1, fqmul(a0, b1) + fqmul(a1, b0))

        // Second pair (uses negative zeta)
        let a2 = t_index_get(a, j + 2)
        let a3 = t_index_get(a, j + 3)
        let b2 = t_index_get(b, j + 2)
        let b3 = t_index_get(b, j + 3)
        let neg_zeta = 0 - zeta
        r = t_index_set(r, j + 2, fqmul(a2, b2) + fqmul(fqmul(a3, b3), neg_zeta))
        r = t_index_set(r, j + 3, fqmul(a2, b3) + fqmul(a3, b2))

        i = i + 1
    }
    return r
}

// ── Self-test ────────────────────────────────────────────────────
// NTT then inverse NTT should return the original polynomial.

fn ntt_self_test() -> Int {
    // Create a test polynomial: [0, 1, 2, ..., 255]
    // The Kyber NTT operates in Montgomery form internally. The round-trip
    // NTT^{-1}(NTT(x)) produces x * R (one extra Montgomery factor).
    // So we convert each output coefficient via fq_from_mont to strip R.
    let mut poly = t_zeros(256)
    let mut i = 0
    while i < 256 {
        poly = t_index_set(poly, i, i + 0.0)
        i = i + 1
    }

    // Forward then inverse
    let ntt_poly = ntt_forward(poly)
    let recovered = ntt_inverse(ntt_poly)

    // Check each coefficient after stripping the Montgomery factor
    let mut j = 0
    let mut errors = 0
    while j < 256 {
        let orig = j % fq_q()
        let raw = t_index_get(recovered, j)
        // Strip Montgomery factor and normalize to [0, q)
        let got = fq_from_mont(raw)
        if got != orig {
            if errors < 5 {
                println("  MISMATCH at " + int_to_str(j) + ": expected " + int_to_str(orig) + " got " + int_to_str(got))
            }
            errors = errors + 1
        }
        j = j + 1
    }
    return errors
}
