plat-core 0.1.1

//! Number Theoretic Transform (NTT) for fast polynomial multiplication
//! in the negacyclic ring Z_q[X]/(X^N + 1).
//!
//! Strategy: twisted NTT.
//!   Forward: twist coefficients by ψ^i, then standard DFT with ω = ψ^2.
//!   Inverse: standard IDFT, then untwist by ψ^{-i}, normalize by N^{-1}.
//!   Pointwise multiply in evaluation domain gives negacyclic convolution.

use crate::modular::{mod_inv, mod_mul, mod_sub};
use crate::params::Params;
use crate::poly::Poly;

/// Precomputed NTT tables.
#[derive(Clone, Debug)]
pub struct NttTables {
    /// ψ^i for i = 0..N (negacyclic twist)
    psi_powers: Vec<u64>,
    /// ψ^{-i} for i = 0..N (inverse twist)
    psi_inv_powers: Vec<u64>,
    n_inv: u64,
    q: u64,
    /// Ring dimension (kept for API consumers).
    pub n: usize,
}

impl NttTables {
    pub fn new(params: &Params) -> Self {
        let n = params.n;
        let q = params.q;
        let psi = params.ntt_root();
        let _omega = mod_mul(psi, psi, q); // N-th root of unity (used in ct/gs_butterfly)

        // Precompute psi powers
        let mut psi_powers = vec![1u64; n];
        for i in 1..n {
            psi_powers[i] = mod_mul(psi_powers[i - 1], psi, q);
        }

        let psi_inv = mod_inv(psi, q);
        let mut psi_inv_powers = vec![1u64; n];
        for i in 1..n {
            psi_inv_powers[i] = mod_mul(psi_inv_powers[i - 1], psi_inv, q);
        }

        let n_inv = mod_inv(n as u64, q);

        Self {
            psi_powers,
            psi_inv_powers,
            n_inv,
            q,
            n,
        }
    }

    /// Forward NTT (coefficient → evaluation domain).
    pub fn forward(&self, poly: &Poly) -> Vec<u64> {
        let q = self.q;

        // Twist: a_i ← a_i * ψ^i
        let mut a: Vec<u64> = poly
            .coeffs
            .iter()
            .enumerate()
            .map(|(i, &c)| mod_mul(c, self.psi_powers[i], q))
            .collect();

        ct_butterfly(&mut a, q, &self.psi_powers);
        a
    }

    /// Inverse NTT (evaluation domain → coefficient).
    pub fn inverse(&self, ntt_values: &[u64]) -> Poly {
        let q = self.q;
        let mut a = ntt_values.to_vec();

        gs_butterfly(&mut a, q, &self.psi_powers);

        // Normalize by N^{-1} and untwist: a_i ← a_i * N^{-1} * ψ^{-i}
        let coeffs = a
            .iter()
            .enumerate()
            .map(|(i, &c)| {
                let c = mod_mul(c, self.n_inv, q);
                mod_mul(c, self.psi_inv_powers[i], q)
            })
            .collect();

        Poly { coeffs }
    }

    /// Pointwise multiply in NTT domain.
    pub fn pointwise_mul(&self, a: &[u64], b: &[u64]) -> Vec<u64> {
        a.iter()
            .zip(b)
            .map(|(&x, &y)| mod_mul(x, y, self.q))
            .collect()
    }

    /// Full polynomial multiplication: a * b in Z_q[X]/(X^N + 1).
    pub fn mul(&self, a: &Poly, b: &Poly) -> Poly {
        let a_ntt = self.forward(a);
        let b_ntt = self.forward(b);
        let c_ntt = self.pointwise_mul(&a_ntt, &b_ntt);
        self.inverse(&c_ntt)
    }
}

/// Cooley-Tukey butterfly (DIT). Uses ω = ψ^2 as N-th root of unity.
/// psi_powers[1] = ψ, so ω = psi_powers[1]^2 mod q. We recompute roots directly.
fn ct_butterfly(a: &mut [u64], q: u64, psi_powers: &[u64]) {
    let n = a.len();
    let omega = mod_mul(psi_powers[1], psi_powers[1], q);

    // Bit-reverse the input
    let log_n = n.trailing_zeros();
    for i in 0..n {
        let j = bit_reverse_idx(i, log_n);
        if i < j {
            a.swap(i, j);
        }
    }

    // Butterfly stages
    let mut len = 2;
    while len <= n {
        let half = len / 2;
        // The twiddle factor step: ω^{N/len}
        let w_step = pow_mod(omega, (n / len) as u64, q);
        for start in (0..n).step_by(len) {
            let mut w = 1u64;
            for j in 0..half {
                let u = a[start + j];
                let v = mod_mul(a[start + j + half], w, q);
                a[start + j] = (u as u128 + v as u128).wrapping_rem(q as u128) as u64;
                a[start + j + half] = mod_sub(u, v, q);
                w = mod_mul(w, w_step, q);
            }
        }
        len *= 2;
    }
}

/// Gentleman-Sande butterfly (DIF) for inverse NTT.
fn gs_butterfly(a: &mut [u64], q: u64, psi_powers: &[u64]) {
    let n = a.len();
    let omega_inv = mod_inv(mod_mul(psi_powers[1], psi_powers[1], q), q);

    let mut len = n;
    while len >= 2 {
        let half = len / 2;
        let w_step = pow_mod(omega_inv, (n / len) as u64, q);
        for start in (0..n).step_by(len) {
            let mut w = 1u64;
            for j in 0..half {
                let u = a[start + j];
                let v = a[start + j + half];
                a[start + j] = (u as u128 + v as u128).wrapping_rem(q as u128) as u64;
                a[start + j + half] = mod_mul(mod_sub(u, v, q), w, q);
                w = mod_mul(w, w_step, q);
            }
        }
        len /= 2;
    }

    // Bit-reverse the output
    let log_n = n.trailing_zeros();
    for i in 0..n {
        let j = bit_reverse_idx(i, log_n);
        if i < j {
            a.swap(i, j);
        }
    }
}

/// Modular exponentiation (small inline version).
fn pow_mod(base: u64, mut exp: u64, modulus: u64) -> u64 {
    let mut result = 1u128;
    let m = modulus as u128;
    let mut b = base as u128;
    while exp > 0 {
        if exp & 1 == 1 {
            result = (result * b) % m;
        }
        exp >>= 1;
        b = (b * b) % m;
    }
    result as u64
}

/// Bit-reverse an index.
fn bit_reverse_idx(i: usize, bits: u32) -> usize {
    i.reverse_bits() >> (usize::BITS - bits)
}

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

    #[test]
    fn test_ntt_roundtrip_tiny() {
        let p = Params::test_tiny();
        let tables = NttTables::new(&p);
        let a = Poly::from_coeffs(&[1, 2, 3, 4, 5], p.n);

        let a_ntt = tables.forward(&a);
        let a_back = tables.inverse(&a_ntt);

        for i in 0..p.n {
            assert_eq!(a.coeffs[i], a_back.coeffs[i], "roundtrip failed at {i}");
        }
    }

    #[test]
    fn test_ntt_roundtrip_random() {
        let p = Params::test_tiny();
        let tables = NttTables::new(&p);

        use rand::SeedableRng;
        use rand::rngs::StdRng;
        let mut rng = StdRng::seed_from_u64(99);
        let a = crate::sampling::sample_uniform(&p, &mut rng);

        let a_ntt = tables.forward(&a);
        let a_back = tables.inverse(&a_ntt);

        for i in 0..p.n {
            assert_eq!(a.coeffs[i], a_back.coeffs[i], "roundtrip failed at {i}");
        }
    }

    #[test]
    fn test_ntt_mul_matches_schoolbook() {
        let p = Params::test_tiny();
        let tables = NttTables::new(&p);
        let a = Poly::from_coeffs(&[1, 2, 3], p.n);
        let b = Poly::from_coeffs(&[4, 5, 6], p.n);

        let c_school = a.mul_schoolbook(&b, &p);
        let c_ntt = tables.mul(&a, &b);

        for i in 0..p.n {
            assert_eq!(
                c_school.coeffs[i], c_ntt.coeffs[i],
                "NTT mul differs from schoolbook at index {i}"
            );
        }
    }

    #[test]
    fn test_ntt_mul_small_params() {
        let p = Params::test_small();
        let tables = NttTables::new(&p);
        let a = Poly::from_coeffs(&[100, 200, 300, 400], p.n);
        let b = Poly::from_coeffs(&[10, 20, 30], p.n);

        let c_school = a.mul_schoolbook(&b, &p);
        let c_ntt = tables.mul(&a, &b);

        for i in 0..p.n {
            assert_eq!(
                c_school.coeffs[i], c_ntt.coeffs[i],
                "NTT mul differs at index {i}"
            );
        }
    }

    #[test]
    fn test_ntt_mul_random() {
        let p = Params::test_small();
        let tables = NttTables::new(&p);

        use rand::SeedableRng;
        use rand::rngs::StdRng;
        let mut rng = StdRng::seed_from_u64(77);
        let a = crate::sampling::sample_uniform(&p, &mut rng);
        let b = crate::sampling::sample_uniform(&p, &mut rng);

        let c_school = a.mul_schoolbook(&b, &p);
        let c_ntt = tables.mul(&a, &b);

        for i in 0..p.n {
            assert_eq!(c_school.coeffs[i], c_ntt.coeffs[i], "differs at {i}");
        }
    }
}