package ring

// Number-Theoretic Transform for R_q = Z_q[x]/(x^n + 1).
// Cooley-Tukey forward, Gentleman-Sande inverse, precomputed twiddle factors.

type nttTables struct {
	n           int32
	q           uint32
	logN        int32
	psiPow      []uint32
	psiInvPow   []uint32
	omegaPow    []uint32
	omegaInvPow []uint32
	bitrevPerm  []int32
	invN        uint32
}

var tableCache map[[2]uint32]*nttTables

func main() {
	tableCache = map[[2]uint32]*nttTables{}
}

func getTables(p Params) (t *nttTables) {
	key := [2]uint32{uint32(p.N), p.Q}
	t, ok := tableCache[key]
	if ok {
		return t
	}
	t = newNTTTables(p)
	tableCache[key] = t
	return t
}

func newNTTTables(p Params) (t *nttTables) {
	n := p.N
	q := p.Q
	psi := p.RootOfUnity
	logN := log2(n)

	psiPow := []uint32{:2 * n}
	psiPow[0] = 1
	for i := int32(1); i < 2*n; i++ {
		psiPow[i] = mulMod(psiPow[i-1], psi, q)
	}

	psiInv := powMod(psi, q-2, q)
	psiInvPow := []uint32{:2 * n}
	psiInvPow[0] = 1
	for i := int32(1); i < 2*n; i++ {
		psiInvPow[i] = mulMod(psiInvPow[i-1], psiInv, q)
	}

	omega := mulMod(psi, psi, q)
	omegaPow := []uint32{:n}
	omegaPow[0] = 1
	for i := int32(1); i < n; i++ {
		omegaPow[i] = mulMod(omegaPow[i-1], omega, q)
	}

	omegaInv := powMod(omega, q-2, q)
	omegaInvPow := []uint32{:n}
	omegaInvPow[0] = 1
	for i := int32(1); i < n; i++ {
		omegaInvPow[i] = mulMod(omegaInvPow[i-1], omegaInv, q)
	}

	bitrevPerm := []int32{:n}
	for i := int32(0); i < n; i++ {
		bitrevPerm[i] = bitrev(i, logN)
	}

	return &nttTables{
		n:           n,
		q:           q,
		logN:        logN,
		psiPow:      psiPow,
		psiInvPow:   psiInvPow,
		omegaPow:    omegaPow,
		omegaInvPow: omegaInvPow,
		bitrevPerm:  bitrevPerm,
		invN:        powMod(uint32(n), q-2, q),
	}
}

func NTT(a *Poly) {
	if a.isNTT {
		return
	}
	t := getTables(a.params)
	n := t.n
	q := t.q
	c := a.Coeffs

	for i := int32(0); i < n; i++ {
		c[i] = mulMod(c[i], t.psiPow[i], q)
	}

	for i := int32(0); i < n; i++ {
		j := t.bitrevPerm[i]
		if i < j {
			c[i], c[j] = c[j], c[i]
		}
	}

	for length := int32(1); length < n; length <<= 1 {
		step := n / (2 * length)
		for start := int32(0); start < n; start += 2 * length {
			for j := int32(0); j < length; j++ {
				tw := t.omegaPow[(j*step)%n]
				idx0 := start + j
				idx1 := idx0 + length
				u := c[idx0]
				v := mulMod(c[idx1], tw, q)
				c[idx0] = addMod(u, v, q)
				c[idx1] = subMod(u, v, q)
			}
		}
	}

	a.isNTT = true
}

func INTT(a *Poly) {
	if !a.isNTT {
		return
	}
	t := getTables(a.params)
	n := t.n
	q := t.q
	c := a.Coeffs

	for i := int32(0); i < n; i++ {
		j := t.bitrevPerm[i]
		if i < j {
			c[i], c[j] = c[j], c[i]
		}
	}

	for length := int32(1); length < n; length <<= 1 {
		step := n / (2 * length)
		for start := int32(0); start < n; start += 2 * length {
			for j := int32(0); j < length; j++ {
				tw := t.omegaInvPow[(j*step)%n]
				idx0 := start + j
				idx1 := idx0 + length
				u := c[idx0]
				v := mulMod(c[idx1], tw, q)
				c[idx0] = addMod(u, v, q)
				c[idx1] = subMod(u, v, q)
			}
		}
	}

	for i := int32(0); i < n; i++ {
		c[i] = mulMod(c[i], t.invN, q)
		c[i] = mulMod(c[i], t.psiInvPow[i], q)
	}

	a.isNTT = false
}

func Mul(a, b *Poly) (c *Poly) {
	if a.isNTT && b.isNTT {
		return MulPointwise(a, b)
	}

	aNTT := a.Clone()
	bNTT := b.Clone()
	NTT(aNTT)
	NTT(bNTT)
	c = MulPointwise(aNTT, bNTT)
	INTT(c)
	return c
}

func powMod(base, exp, q uint32) (result uint32) {
	result = 1
	b := base % q
	for e := exp; e > 0; e >>= 1 {
		if e&1 == 1 {
			result = mulMod(result, b, q)
		}
		b = mulMod(b, b, q)
	}
	return result
}

func log2(n int32) (r int32) {
	n >>= 1
	for n > 0 {
		r++
		n >>= 1
	}
	return r
}

func bitrev(x, bits int32) (r int32) {
	for i := int32(0); i < bits; i++ {
		r = (r << 1) | (x & 1)
		x >>= 1
	}
	return r
}