use crypto::fieldq
use crypto::ntt
use crypto::sample
use crypto::kpke

fn main() {
    println("=== K-PKE k=1 NO NOISE, message only ===")
    let rho = secure_random_bytes(32)
    let sigma = secure_random_bytes(32)

    // Keygen: t = A*s (NO error)
    let a = sample_ntt(rho, 0, 0)
    let s_ntt = ntt_forward(sample_noise_cbd2(sigma, 0))
    let t_ntt = ntt_basemul(a, s_ntt)

    // Encrypt with r, NO noise (e1=0, e2=0)
    let r_seed = secure_random_bytes(32)
    let r_ntt = ntt_forward(sample_noise_cbd2(r_seed, 0))

    let u_ntt = ntt_basemul(a, r_ntt)
    let u = ntt_inverse(u_ntt)

    let v_ntt = ntt_basemul(t_ntt, r_ntt)
    let v_base = ntt_inverse(v_ntt)

    // Add message: bit 1 at every even position
    let mut msg_poly = t_zeros(256)
    let mut mi = 0
    while mi < 256 {
        if mi % 2 == 0 {
            msg_poly = t_index_set(msg_poly, mi, 1665.0)
        }
        mi = mi + 1
    }
    let v = poly_add(v_base, msg_poly)

    // Decrypt: su = INTT(basemul(s, NTT(u)))
    let su_ntt = ntt_basemul(s_ntt, ntt_forward(u))
    let su = ntt_inverse(su_ntt)

    // w = v - su should be msg_poly exactly (no noise)
    let w = poly_sub(v, su)

    println("First 8 recovered coefficients (expect 1665, 0, 1665, 0, ...):")
    let mut j = 0
    while j < 8 {
        let coeff = fq_normalize(t_index_get(w, j))
        println("  [" + int_to_str(j) + "] = " + int_to_str(coeff))
        j = j + 1
    }

    // BUT ALSO: what if NTT(INTT(x)) != x? Test this:
    println("")
    println("NTT(INTT(u_ntt)) vs u_ntt (first 4):")
    let u_roundtrip = ntt_forward(ntt_inverse(u_ntt))
    let mut k = 0
    while k < 4 {
        let orig = t_index_get(u_ntt, k)
        let rt = t_index_get(u_roundtrip, k)
        println("  [" + int_to_str(k) + "] orig=" + float_to_str(orig) + " roundtrip=" + float_to_str(rt) + " ratio=" + float_to_str(rt / orig))
        k = k + 1
    }
}
