package crypto

// Number-Theoretic Transform (NTT) over Z_257 for Hamadryad.
//
// The NTT is the finite-field analogue of the FFT. It transforms
// a polynomial from coefficient form to evaluation form, enabling
// O(n log n) polynomial multiplication via pointwise multiply.
//
// For Hamadryad: n=64, p=257.
//
// 257 is prime and 256 = 257-1 = 2^8, so Z_257* has order 256.
// g = 3 is a primitive root (generator) of Z_257*.
// psi = 3^(256/128) = 3^2 = 9 is a primitive 128th root of unity.
//
// For polynomials mod (x^64 + 1), we use the negacyclic NTT:
// pre-multiply coefficients by powers of psi (a 128th root),
// then apply a standard 64-point NTT with omega = psi^2 = 81
// (a primitive 64th root of unity).

// mod257 reduces x into [0, 256].
func mod257(x int) uint16 {
	x %= HamP
	if x < 0 {
		x += HamP
	}
	return uint16(x)
}

// powMod computes base^exp mod 257.
func powMod(base, exp int) uint16 {
	result := 1
	b := base % HamP
	if b < 0 {
		b += HamP
	}
	for e := exp; e > 0; e >>= 1 {
		if e&1 == 1 {
			result = (result * b) % HamP
		}
		b = (b * b) % HamP
	}
	return uint16(result)
}

// invMod computes the modular inverse of a mod 257.
// Uses Fermat's little theorem: a^(-1) = a^(p-2) mod p.
func invMod(a uint16) uint16 {
	return powMod(int(a), HamP-2)
}

// bitRev6 reverses the low 6 bits of x.
func bitRev6(x int) int {
	r := 0
	for range 6 {
		r = (r << 1) | (x & 1)
		x >>= 1
	}
	return r
}

// Pre-computed tables, filled by initNTTTables.
var (
	// psiPows[i] = psi^i mod 257, for i=0..127.
	// psi = 9 (primitive 128th root of unity).
	psiPows [128]uint16

	// psiInvPows[i] = psi^(-i) mod 257.
	psiInvPows [128]uint16

	// invN = 64^(-1) mod 257.
	invN uint16

	nttTablesReady bool
)

// initNTTTables populates the NTT lookup tables. Safe to call multiple
// times; only the first call has effect. Called explicitly from the
// hamadryad init to guarantee table availability before key generation
// (Go runs init functions in source file alphabetical order within a
// package, and "hamadryad.go" sorts before "ntt.go").
func initNTTTables() {
	if nttTablesReady {
		return
	}
	const psi = 9 // primitive 128th root of unity in Z_257

	psiPows[0] = 1
	for i := 1; i < 128; i++ {
		psiPows[i] = mod257(int(psiPows[i-1]) * psi)
	}

	psiInv := invMod(psi)
	psiInvPows[0] = 1
	for i := 1; i < 128; i++ {
		psiInvPows[i] = mod257(int(psiInvPows[i-1]) * int(psiInv))
	}

	invN = invMod(uint16(HamN))
	nttTablesReady = true
}

func init() {
	initNTTTables()
}

// ntt64 computes the forward negacyclic NTT of a length-64 polynomial over Z_257.
//
// The negacyclic NTT computes polynomial evaluation at the roots of (x^64 + 1),
// which are psi^(2k+1) for k=0..63. Implementation:
//  1. Pre-multiply: a[i] *= psi^i
//  2. Standard radix-2 DIT NTT with omega = psi^2 (primitive 64th root)
func ntt64(a *[HamN]uint16) {
	const n = HamN

	// Step 1: pre-multiply by psi^i.
	for i := range n {
		a[i] = mod257(int(a[i]) * int(psiPows[i]))
	}

	// Step 2: bit-reversal permutation.
	for i := range n {
		j := bitRev6(i)
		if i < j {
			a[i], a[j] = a[j], a[i]
		}
	}

	// Step 3: Cooley-Tukey butterfly stages.
	// omega = psi^2 = 9^2 = 81 (primitive 64th root of unity).
	for length := 1; length < n; length <<= 1 {
		// Twiddle factor base: omega^(n/(2*length)).
		step := n / (2 * length)
		for start := 0; start < n; start += 2 * length {
			for j := range length {
				// tw = omega^(j*step) = psi^(2*j*step)
				tw := psiPows[(2*j*step)%128]
				u := a[start+j]
				v := mod257(int(a[start+j+length]) * int(tw))
				a[start+j] = mod257(int(u) + int(v))
				a[start+j+length] = mod257(int(u) - int(v))
			}
		}
	}
}

// intt64 computes the inverse negacyclic NTT, recovering coefficients.
func intt64(a *[HamN]uint16) {
	const n = HamN

	// Step 1: bit-reversal permutation.
	for i := range n {
		j := bitRev6(i)
		if i < j {
			a[i], a[j] = a[j], a[i]
		}
	}

	// Step 2: Gentleman-Sande butterfly (inverse DIT) with omega^(-1).
	for length := 1; length < n; length <<= 1 {
		step := n / (2 * length)
		for start := 0; start < n; start += 2 * length {
			for j := range length {
				// tw = omega^(-(j*step)) = psi^(-2*j*step)
				tw := psiInvPows[(2*j*step)%128]
				u := a[start+j]
				v := mod257(int(a[start+j+length]) * int(tw))
				a[start+j] = mod257(int(u) + int(v))
				a[start+j+length] = mod257(int(u) - int(v))
			}
		}
	}

	// Step 3: multiply by 1/n.
	for i := range n {
		a[i] = mod257(int(a[i]) * int(invN))
	}

	// Step 4: undo psi pre-multiplication: a[i] *= psi^(-i).
	for i := range n {
		a[i] = mod257(int(a[i]) * int(psiInvPows[i]))
	}
}