package hamcrypto

// Radix-3 NTT over Z_271 for trinary Hamadryad.
// n=27 = 3^3, p=271. psi=188 (primitive 54th root of unity).

const GnarlP = 271
const GnarlN = 27

// mod271 reduces x into [0, 270] using Barrett reduction.
func mod271(x uint32) (r uint16) {
	q := (x * 483) >> 17
	rv := x - q*271
	if rv >= 271 {
		rv -= 271
	}
	return uint16(rv)
}

// mod271s reduces a signed int into [0, 270].
func mod271s(x int32) (r uint16) {
	xv := x % GnarlP
	if xv < 0 {
		xv += GnarlP
	}
	return uint16(xv)
}

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

// invMod271 computes the modular inverse of a mod 271.
func invMod271(a uint16) (r uint16) {
	return powMod271(int32(a), GnarlP-2)
}

// digitRev3 reverses the base-3 digits of x with the given number of digits.
func digitRev3(x, digits int32) (r int32) {
	r = 0
	xv := x
	for i := int32(0); i < digits; i++ {
		r = r*3 + xv%3
		xv /= 3
	}
	return r
}

// Pre-computed tables for the radix-3 NTT.
var psiPows27 [54]uint16
var psiInvPows27 [54]uint16
var invN27 uint16
var ntt27TablesReady bool

// initNTT27Tables populates the radix-3 NTT lookup tables.
func initNTT27Tables() {
	if ntt27TablesReady {
		return
	}
	const psi = 188

	psiPows27[0] = 1
	for i := int32(1); i < 54; i++ {
		psiPows27[i] = mod271(uint32(psiPows27[i-1]) * psi)
	}

	psiInv := invMod271(psi)
	psiInvPows27[0] = 1
	for i := int32(1); i < 54; i++ {
		psiInvPows27[i] = mod271(uint32(psiInvPows27[i-1]) * uint32(psiInv))
	}

	invN27 = invMod271(uint16(GnarlN))
	ntt27TablesReady = true
}

// ntt27 computes the forward negacyclic NTT of a length-27 polynomial over Z_271.
func ntt27(a *[GnarlN]uint16) {
	// Pre-multiply by psi^i for negacyclic.
	a[1] = mod271(uint32(a[1]) * 188)
	a[2] = mod271(uint32(a[2]) * 114)
	a[3] = mod271(uint32(a[3]) * 23)
	a[4] = mod271(uint32(a[4]) * 259)
	a[5] = mod271(uint32(a[5]) * 183)
	a[6] = mod271(uint32(a[6]) * 258)
	a[7] = mod271(uint32(a[7]) * 266)
	a[8] = mod271(uint32(a[8]) * 144)
	a[9] = mod271(uint32(a[9]) * 243)
	a[10] = mod271(uint32(a[10]) * 156)
	a[11] = mod271(uint32(a[11]) * 60)
	a[12] = mod271(uint32(a[12]) * 169)
	a[13] = mod271(uint32(a[13]) * 65)
	a[14] = mod271(uint32(a[14]) * 25)
	a[15] = mod271(uint32(a[15]) * 93)
	a[16] = mod271(uint32(a[16]) * 140)
	a[17] = mod271(uint32(a[17]) * 33)
	a[18] = mod271(uint32(a[18]) * 242)
	a[19] = mod271(uint32(a[19]) * 239)
	a[20] = mod271(uint32(a[20]) * 217)
	a[21] = mod271(uint32(a[21]) * 146)
	a[22] = mod271(uint32(a[22]) * 77)
	a[23] = mod271(uint32(a[23]) * 113)
	a[24] = mod271(uint32(a[24]) * 106)
	a[25] = mod271(uint32(a[25]) * 145)
	a[26] = mod271(uint32(a[26]) * 160)

	// Digit-reversal permutation (9 swaps).
	a[1], a[9] = a[9], a[1]
	a[2], a[18] = a[18], a[2]
	a[4], a[12] = a[12], a[4]
	a[5], a[21] = a[21], a[5]
	a[7], a[15] = a[15], a[7]
	a[8], a[24] = a[24], a[8]
	a[11], a[19] = a[19], a[11]
	a[14], a[22] = a[22], a[14]
	a[17], a[25] = a[25], a[17]

	// Stage 0: stride=1, all tw1=tw2=1
	nttBfly1(a, 0, 1, 2)
	nttBfly1(a, 3, 4, 5)
	nttBfly1(a, 6, 7, 8)
	nttBfly1(a, 9, 10, 11)
	nttBfly1(a, 12, 13, 14)
	nttBfly1(a, 15, 16, 17)
	nttBfly1(a, 18, 19, 20)
	nttBfly1(a, 21, 22, 23)
	nttBfly1(a, 24, 25, 26)

	// Stage 1: stride=3
	nttBfly1(a, 0, 3, 6)
	nttBfly1(a, 9, 12, 15)
	nttBfly1(a, 18, 21, 24)
	nttBfly(a, 1, 4, 7, 258, 169)
	nttBfly(a, 10, 13, 16, 258, 169)
	nttBfly(a, 19, 22, 25, 258, 169)
	nttBfly(a, 2, 5, 8, 169, 106)
	nttBfly(a, 11, 14, 17, 169, 106)
	nttBfly(a, 20, 23, 26, 169, 106)

	// Stage 2: stride=9
	nttBfly1(a, 0, 9, 18)
	nttBfly(a, 1, 10, 19, 114, 259)
	nttBfly(a, 2, 11, 20, 259, 144)
	nttBfly(a, 3, 12, 21, 258, 169)
	nttBfly(a, 4, 13, 22, 144, 140)
	nttBfly(a, 5, 14, 23, 156, 217)
	nttBfly(a, 6, 15, 24, 169, 106)
	nttBfly(a, 7, 16, 25, 25, 83)
	nttBfly(a, 8, 17, 26, 140, 88)
}

// nttBfly1 performs a radix-3 butterfly with trivial twiddles.
func nttBfly1(a *[27]uint16, i0, i1, i2 int32) {
	v0 := uint32(a[i0])
	v1 := uint32(a[i1])
	v2 := uint32(a[i2])
	a[i0] = mod271(v0 + v1 + v2)
	a[i1] = mod271(v0 + uint32(mod271(v1*242)) + uint32(mod271(v2*28)))
	a[i2] = mod271(v0 + uint32(mod271(v1*28)) + uint32(mod271(v2*242)))
}

// nttBfly performs a radix-3 butterfly with given twiddle factors.
func nttBfly(a *[27]uint16, i0, i1, i2 int32, tw1, tw2 uint32) {
	v0 := uint32(a[i0])
	a1tw := uint32(mod271(uint32(a[i1]) * tw1))
	a2tw := uint32(mod271(uint32(a[i2]) * tw2))
	a[i0] = mod271(v0 + a1tw + a2tw)
	a[i1] = mod271(v0 + uint32(mod271(a1tw*242)) + uint32(mod271(a2tw*28)))
	a[i2] = mod271(v0 + uint32(mod271(a1tw*28)) + uint32(mod271(a2tw*242)))
}

// intt27 computes the inverse negacyclic NTT, recovering coefficients.
func intt27(a *[GnarlN]uint16) {
	// Stage 2 (DIF, top-down)
	inttBfly1(a, 0, 9, 18)
	inttBfly(a, 1, 10, 19, 126, 158)
	inttBfly(a, 2, 11, 20, 158, 32)
	inttBfly(a, 3, 12, 21, 125, 178)
	inttBfly(a, 4, 13, 22, 32, 211)
	inttBfly(a, 5, 14, 23, 238, 5)
	inttBfly(a, 6, 15, 24, 178, 248)
	inttBfly(a, 7, 16, 25, 206, 160)
	inttBfly(a, 8, 17, 26, 211, 77)

	// Stage 1
	inttBfly1(a, 0, 3, 6)
	inttBfly(a, 1, 4, 7, 125, 178)
	inttBfly(a, 2, 5, 8, 178, 248)
	inttBfly1(a, 9, 12, 15)
	inttBfly(a, 10, 13, 16, 125, 178)
	inttBfly(a, 11, 14, 17, 178, 248)
	inttBfly1(a, 18, 21, 24)
	inttBfly(a, 19, 22, 25, 125, 178)
	inttBfly(a, 20, 23, 26, 178, 248)

	// Stage 0
	inttBfly1(a, 0, 1, 2)
	inttBfly1(a, 3, 4, 5)
	inttBfly1(a, 6, 7, 8)
	inttBfly1(a, 9, 10, 11)
	inttBfly1(a, 12, 13, 14)
	inttBfly1(a, 15, 16, 17)
	inttBfly1(a, 18, 19, 20)
	inttBfly1(a, 21, 22, 23)
	inttBfly1(a, 24, 25, 26)

	// Digit-reversal permutation.
	a[1], a[9] = a[9], a[1]
	a[2], a[18] = a[18], a[2]
	a[4], a[12] = a[12], a[4]
	a[5], a[21] = a[21], a[5]
	a[7], a[15] = a[15], a[7]
	a[8], a[24] = a[24], a[8]
	a[11], a[19] = a[19], a[11]
	a[14], a[22] = a[22], a[14]
	a[17], a[25] = a[25], a[17]

	// Post-multiply: fused invN27 * psiInvPows27[i].
	a[0] = mod271(uint32(a[0]) * 261)
	a[1] = mod271(uint32(a[1]) * 245)
	a[2] = mod271(uint32(a[2]) * 95)
	a[3] = mod271(uint32(a[3]) * 247)
	a[4] = mod271(uint32(a[4]) * 46)
	a[5] = mod271(uint32(a[5]) * 228)
	a[6] = mod271(uint32(a[6]) * 105)
	a[7] = mod271(uint32(a[7]) * 2)
	a[8] = mod271(uint32(a[8]) * 222)
	a[9] = mod271(uint32(a[9]) * 252)
	a[10] = mod271(uint32(a[10]) * 59)
	a[11] = mod271(uint32(a[11]) * 45)
	a[12] = mod271(uint32(a[12]) * 117)
	a[13] = mod271(uint32(a[13]) * 250)
	a[14] = mod271(uint32(a[14]) * 108)
	a[15] = mod271(uint32(a[15]) * 64)
	a[16] = mod271(uint32(a[16]) * 58)
	a[17] = mod271(uint32(a[17]) * 205)
	a[18] = mod271(uint32(a[18]) * 262)
	a[19] = mod271(uint32(a[19]) * 85)
	a[20] = mod271(uint32(a[20]) * 221)
	a[21] = mod271(uint32(a[21]) * 141)
	a[22] = mod271(uint32(a[22]) * 204)
	a[23] = mod271(uint32(a[23]) * 151)
	a[24] = mod271(uint32(a[24]) * 230)
	a[25] = mod271(uint32(a[25]) * 56)
	a[26] = mod271(uint32(a[26]) * 254)
}

// inttBfly1 performs an inverse radix-3 DIF butterfly with trivial twiddles.
func inttBfly1(a *[27]uint16, i0, i1, i2 int32) {
	v0 := uint32(a[i0])
	v1 := uint32(a[i1])
	v2 := uint32(a[i2])
	a[i0] = mod271(v0 + v1 + v2)
	a[i1] = mod271(v0 + uint32(mod271(v1*28)) + uint32(mod271(v2*242)))
	a[i2] = mod271(v0 + uint32(mod271(v1*242)) + uint32(mod271(v2*28)))
}

// inttBfly performs an inverse radix-3 DIF butterfly with given twiddle factors.
func inttBfly(a *[27]uint16, i0, i1, i2 int32, tw1inv, tw2inv uint32) {
	v0 := uint32(a[i0])
	v1 := uint32(a[i1])
	v2 := uint32(a[i2])
	a[i0] = mod271(v0 + v1 + v2)
	b1 := mod271(v0 + uint32(mod271(v1*28)) + uint32(mod271(v2*242)))
	b2 := mod271(v0 + uint32(mod271(v1*242)) + uint32(mod271(v2*28)))
	a[i1] = mod271(uint32(b1) * tw1inv)
	a[i2] = mod271(uint32(b2) * tw2inv)
}