// Package ring implements polynomial arithmetic over R_q = Z_q[x]/(x^n + 1).
//
// The ring R_q = Z_q[x]/(x^n + 1) is a quotient of the polynomial ring Z_q[x]
// by the cyclotomic polynomial Phi_{2n}(x) = x^n + 1. When n is a power of 2
// and q = 1 (mod 2n), this ring supports a negacyclic NTT, enabling O(n log n)
// polynomial multiplication.
package ring

type Params struct {
	N           int32
	Q           uint32
	RootOfUnity uint32
	MontR       uint64
	QInv        uint64
}

func Falcon512() (p Params) {
	return Params{
		N:           512,
		Q:           12289,
		RootOfUnity: 49,
		MontR:       1 << 16,
		QInv:        qinv(12289, 16),
	}
}

func NewHope256() (p Params) {
	return Params{
		N:           256,
		Q:           7681,
		RootOfUnity: 4055,
		MontR:       1 << 16,
		QInv:        qinv(7681, 16),
	}
}

func HE64() (p Params) {
	return Params{
		N:           64,
		Q:           10000769,
		RootOfUnity: 6028202,
		MontR:       1 << 32,
		QInv:        qinv(10000769, 32),
	}
}

func Falcon1024() (p Params) {
	return Params{
		N:           1024,
		Q:           12289,
		RootOfUnity: 1945,
		MontR:       1 << 16,
		QInv:        qinv(12289, 16),
	}
}

func qinv(q uint32, k int32) (result uint64) {
	inv := uint64(1)
	for i := int32(1); i < k; i++ {
		inv = inv * (2 - inv*uint64(q))
	}
	mask := (uint64(1) << uint32(k)) - 1
	return (-inv) & mask
}

type Poly struct {
	Coeffs []uint32
	params Params
	isNTT  bool
}

func New(p Params) (poly *Poly) {
	return &Poly{
		Coeffs: []uint32{:p.N},
		params: p,
	}
}

func NewFromCoeffs(p Params, coeffs []uint32) (poly *Poly) {
	poly = New(p)
	n := int32(len(coeffs))
	if n > p.N {
		n = p.N
	}
	for i := int32(0); i < n; i++ {
		poly.Coeffs[i] = coeffs[i] % p.Q
	}
	return poly
}

func (a *Poly) Clone() (b *Poly) {
	b = &Poly{
		Coeffs: []uint32{:int32(len(a.Coeffs))},
		params: a.params,
		isNTT:  a.isNTT,
	}
	copy(b.Coeffs, a.Coeffs)
	return b
}

func (a *Poly) RingParams() (p Params) { return a.params }
func (a *Poly) IsNTT() (v bool)       { return a.isNTT }

func addMod(a, b, q uint32) (result uint32) {
	sum := a + b
	if sum >= q {
		sum -= q
	}
	return sum
}

func subMod(a, b, q uint32) (result uint32) {
	if a >= b {
		return a - b
	}
	return q - b + a
}

func mulMod(a, b, q uint32) (result uint32) {
	return uint32((uint64(a) * uint64(b)) % uint64(q))
}

func Add(a, b *Poly) (c *Poly) {
	c = New(a.params)
	c.isNTT = a.isNTT
	q := a.params.Q
	for i := int32(0); i < int32(len(a.Coeffs)); i++ {
		c.Coeffs[i] = addMod(a.Coeffs[i], b.Coeffs[i], q)
	}
	return c
}

func Sub(a, b *Poly) (c *Poly) {
	c = New(a.params)
	c.isNTT = a.isNTT
	q := a.params.Q
	for i := int32(0); i < int32(len(a.Coeffs)); i++ {
		c.Coeffs[i] = subMod(a.Coeffs[i], b.Coeffs[i], q)
	}
	return c
}

func Neg(a *Poly) (c *Poly) {
	c = New(a.params)
	c.isNTT = a.isNTT
	q := a.params.Q
	for i := int32(0); i < int32(len(a.Coeffs)); i++ {
		if a.Coeffs[i] != 0 {
			c.Coeffs[i] = q - a.Coeffs[i]
		}
	}
	return c
}

func MulPointwise(a, b *Poly) (c *Poly) {
	c = New(a.params)
	c.isNTT = true
	q := a.params.Q
	for i := int32(0); i < int32(len(a.Coeffs)); i++ {
		c.Coeffs[i] = mulMod(a.Coeffs[i], b.Coeffs[i], q)
	}
	return c
}

func MulAdd(c, a, b *Poly) {
	q := a.params.Q
	for i := int32(0); i < int32(len(a.Coeffs)); i++ {
		c.Coeffs[i] = addMod(c.Coeffs[i], mulMod(a.Coeffs[i], b.Coeffs[i], q), q)
	}
}

func ScalarMul(a *Poly, s uint32) (c *Poly) {
	c = New(a.params)
	c.isNTT = a.isNTT
	q := a.params.Q
	s = s % q
	for i := int32(0); i < int32(len(a.Coeffs)); i++ {
		c.Coeffs[i] = mulMod(a.Coeffs[i], s, q)
	}
	return c
}

func Equal(a, b *Poly) (result bool) {
	if int32(len(a.Coeffs)) != int32(len(b.Coeffs)) {
		return false
	}
	for i := int32(0); i < int32(len(a.Coeffs)); i++ {
		if a.Coeffs[i] != b.Coeffs[i] {
			return false
		}
	}
	return true
}

// LE byte helpers (replacing encoding/binary).
func leU16(b []byte) (v uint16) {
	return uint16(b[0]) | uint16(b[1])<<8
}

func leU32(b []byte) (v uint32) {
	return uint32(b[0]) | uint32(b[1])<<8 | uint32(b[2])<<16 | uint32(b[3])<<24
}

func leU64(b []byte) (v uint64) {
	return uint64(b[0]) | uint64(b[1])<<8 | uint64(b[2])<<16 | uint64(b[3])<<24 |
		uint64(b[4])<<32 | uint64(b[5])<<40 | uint64(b[6])<<48 | uint64(b[7])<<56
}

func putLeU64(b []byte, v uint64) {
	b[0] = byte(v)
	b[1] = byte(v >> 8)
	b[2] = byte(v >> 16)
	b[3] = byte(v >> 24)
	b[4] = byte(v >> 32)
	b[5] = byte(v >> 40)
	b[6] = byte(v >> 48)
	b[7] = byte(v >> 56)
}

func putLeU16(b []byte, v uint16) {
	b[0] = byte(v)
	b[1] = byte(v >> 8)
}

func Serialize(a *Poly) (out []byte) {
	q := a.params.Q
	bits := int32(0)
	for v := q - 1; v > 0; v >>= 1 {
		bits++
	}

	totalBits := bits * int32(len(a.Coeffs))
	out = []byte{:(totalBits + 7) / 8}

	bitPos := int32(0)
	for i := int32(0); i < int32(len(a.Coeffs)); i++ {
		c := a.Coeffs[i]
		for b := int32(0); b < bits; b++ {
			if c&(1<<uint32(b)) != 0 {
				out[bitPos/8] |= 1 << uint32(bitPos%8)
			}
			bitPos++
		}
	}
	return out
}

func Deserialize(p Params, data []byte) (poly *Poly) {
	q := p.Q
	bits := int32(0)
	for v := q - 1; v > 0; v >>= 1 {
		bits++
	}

	poly = New(p)
	bitPos := int32(0)
	for i := int32(0); i < int32(len(poly.Coeffs)); i++ {
		var v uint32
		for b := int32(0); b < bits; b++ {
			if bitPos/8 < int32(len(data)) && data[bitPos/8]&(1<<uint32(bitPos%8)) != 0 {
				v |= 1 << uint32(b)
			}
			bitPos++
		}
		if v >= q {
			v = q - 1
		}
		poly.Coeffs[i] = v
	}
	return poly
}