plat-core 0.1.1

//! Parameter sets for lattice-based FHE.
//!
//! We use NTT-friendly primes: q ≡ 1 (mod 2N) so that primitive 2N-th roots
//! of unity exist in Z_q.

use crate::modular::mod_pow;

/// FHE parameter set.
#[derive(Debug, Clone, Copy)]
pub struct Params {
    /// Ring dimension (must be a power of 2).
    pub n: usize,
    /// Ciphertext modulus (must be prime, q ≡ 1 mod 2N).
    pub q: u64,
    /// Gaussian error standard deviation (as fixed-point × 1000).
    pub sigma_x1000: u64,
    /// Plaintext modulus for encoding.
    pub t: u64,
}

impl Params {
    /// Research-grade parameters: N=2048, ~100-bit security.
    /// q is an NTT-friendly prime: q ≡ 1 (mod 2*2048 = 4096).
    pub fn research_2048() -> Self {
        // q = 2^40 - 2^14 + 1 = 1099511611393 — but let's use a known good one.
        // q = 132120577 (prime, ≡ 1 mod 4096, fits in u64 comfortably for research)
        // Actually for research quality let's use a larger prime for noise headroom.
        // q = 1152921504606830593 (2^60 - 2^14 + 1) — too large for comfortable NTT.
        //
        // Practical choice: q = 0xFFFFFFFF00000001 = 2^64 - 2^32 + 1 (Goldilocks prime)
        // This is NTT-friendly with N up to 2^31.
        // But it's close to u64::MAX, risking overflow in u128 multiplications — fine.
        //
        // Simpler: q = 132120577 = 2^27 - 2^13 + 1, prime, ≡ 1 mod 4096.
        // This gives ~27 bits of modulus, enough for depth-1 circuits with research params.
        Self {
            n: 2048,
            q: 132120577,
            sigma_x1000: 3200, // σ = 3.2
            t: 65537,          // plaintext modulus, prime
        }
    }

    /// Small parameters for fast testing: N=256.
    /// q/t ratio ≈ 48, giving about 5 bits of noise budget per coefficient.
    pub fn test_small() -> Self {
        // q = 12289 is a classic lattice prime, 12289 ≡ 1 mod 512.
        // t = 256: q/t ≈ 48, enough for small noise and a few additions.
        // σ = 1.0 to keep noise small for testing.
        Self {
            n: 256,
            q: 12289,
            sigma_x1000: 1000, // σ = 1.0
            t: 17,             // small plaintext modulus for adequate q/t ratio (~723)
        }
    }

    /// Tiny parameters for unit tests only: N=64.
    pub fn test_tiny() -> Self {
        // q = 769, t = 5, q/t ≈ 153. σ = 1.0.
        // 769 - 1 = 768 = 2^8 * 3. 2N = 128. 768 / 128 = 6 ✓
        Self {
            n: 64,
            q: 769,
            sigma_x1000: 1000, // σ = 1.0
            t: 5,
        }
    }

    /// Find a primitive 2N-th root of unity modulo q.
    pub fn ntt_root(&self) -> u64 {
        let two_n = (2 * self.n) as u64;
        // q - 1 must be divisible by 2N
        assert!(
            (self.q - 1) % two_n == 0,
            "q - 1 = {} is not divisible by 2N = {two_n}",
            self.q - 1
        );
        // Find generator: try small values
        let exp = (self.q - 1) / two_n;
        for g in 2..self.q {
            let w = mod_pow(g, exp, self.q);
            if w != 1 && mod_pow(w, self.n as u64, self.q) != 1 {
                // w is a primitive 2N-th root if w^N ≡ -1 mod q
                let wn = mod_pow(w, self.n as u64, self.q);
                if wn == self.q - 1 {
                    return w;
                }
            }
        }
        panic!("no primitive 2N-th root of unity found");
    }
}

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

    #[test]
    fn test_tiny_root() {
        let p = Params::test_tiny();
        let w = p.ntt_root();
        let wn = mod_pow(w, p.n as u64, p.q);
        assert_eq!(wn, p.q - 1, "w^N should equal -1 mod q");
        let w2n = mod_pow(w, (2 * p.n) as u64, p.q);
        assert_eq!(w2n, 1, "w^(2N) should equal 1 mod q");
    }

    #[test]
    fn test_small_root() {
        let p = Params::test_small();
        let w = p.ntt_root();
        let wn = mod_pow(w, p.n as u64, p.q);
        assert_eq!(wn, p.q - 1);
    }
}