package hamcrypto

// Number-Theoretic Transform (NTT) over Z_257 for Hamadryad.
// n=64, p=257. psi=9 (primitive 128th root of unity).

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

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

// invMod computes the modular inverse of a mod 257.
func invMod(a uint16) (r uint16) {
	return powMod(int32(a), HamP-2)
}

// bitRev6 reverses the low 6 bits of x.
func bitRev6(x int32) (r int32) {
	r = 0
	xv := x
	for i := int32(0); i < 6; i++ {
		r = (r << 1) | (xv & 1)
		xv >>= 1
	}
	return r
}

// Pre-computed tables, filled by initNTTTables.
var psiPows [128]uint16
var psiInvPows [128]uint16
var invN uint16
var nttTablesReady bool

// initNTTTables populates the NTT lookup tables.
func initNTTTables() {
	if nttTablesReady {
		return
	}
	const psi = 9

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

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

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

// ntt64 computes the forward negacyclic NTT of a length-64 polynomial over Z_257.
func ntt64(a *[HamN]uint16) {
	const n = int32(HamN)

	// Pre-multiply by psi^i.
	for i := int32(0); i < n; i++ {
		a[i] = mod257(int32(a[i]) * int32(psiPows[i]))
	}

	// Bit-reversal permutation.
	for i := int32(0); i < n; i++ {
		j := bitRev6(i)
		if i < j {
			a[i], a[j] = a[j], a[i]
		}
	}

	// Cooley-Tukey butterfly stages.
	length := int32(1)
	for length < n {
		step := n / (2 * length)
		start := int32(0)
		for start < n {
			for j := int32(0); j < length; j++ {
				tw := psiPows[(2*j*step)%128]
				u := a[start+j]
				v := mod257(int32(a[start+j+length]) * int32(tw))
				a[start+j] = mod257(int32(u) + int32(v))
				a[start+j+length] = mod257(int32(u) - int32(v))
			}
			start += 2 * length
		}
		length <<= 1
	}
}

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

	// Bit-reversal permutation.
	for i := int32(0); i < n; i++ {
		j := bitRev6(i)
		if i < j {
			a[i], a[j] = a[j], a[i]
		}
	}

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

	// Multiply by 1/n.
	for i := int32(0); i < n; i++ {
		a[i] = mod257(int32(a[i]) * int32(invN))
	}

	// Undo psi pre-multiplication.
	for i := int32(0); i < n; i++ {
		a[i] = mod257(int32(a[i]) * int32(psiInvPows[i]))
	}
}