// Package ring implements polynomial arithmetic over R_q = Z_q[x]/(x^n + 1).
//
// This is the algebraic engine for the hamadryad cryptosystem. All lattice
// operations — hash (SIS), signatures (GPV), encryption (LWE), homomorphic
// evaluation — reduce to polynomial arithmetic in this ring.
//
// 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.
//
// Parameters are selected for cryptographic use:
//
//	Falcon-512:    n=512,  q=12289       (signature scheme, ~128-bit security)
//	Kyber-512:     n=256,  q=3329        (KEM, ~128-bit security)
//	Dilithium-II:  n=256,  q=8380417     (signature, ~128-bit security)
//
// The default parameter set targets Falcon-512 (n=512, q=12289) because:
//   - Hash-and-sign naturally extends hamadryad (both are SIS-based)
//   - n=512 gives 128-bit post-quantum security
//   - q=12289 is NTT-friendly: 12289 ≡ 1 (mod 1024), primitive 1024th root exists
//   - All coefficients fit in uint16 (q < 2^14), enabling fast arithmetic
//
// Security: all constructions reduce to SVP on ideal lattices in R = Z[x]/(x^n+1).
package ring

// Params defines the ring R_q = Z_q[x]/(x^n + 1).
type Params struct {
	// N is the ring dimension (degree of cyclotomic polynomial).
	// Must be a power of 2.
	N int

	// Q is the coefficient modulus. Must be prime with Q ≡ 1 (mod 2N)
	// to guarantee existence of primitive 2N-th roots of unity for NTT.
	Q uint32

	// RootOfUnity is a primitive 2N-th root of unity mod Q.
	// Used for the negacyclic NTT: psi such that psi^(2N) ≡ 1 (mod Q)
	// and psi^N ≡ -1 (mod Q).
	RootOfUnity uint32

	// MontR is the Montgomery parameter R = 2^k where k is chosen
	// so that R > Q. For Q < 2^16, k=16. For Q < 2^32, k=32.
	MontR uint64

	// QInv is -Q^{-1} mod R (Montgomery reduction constant).
	QInv uint64
}

// Falcon512 returns parameters for the Falcon-512 ring.
//
//	R_q = Z_12289[x]/(x^512 + 1)
//	12289 = 12·1024 + 1 = 3·2^12 + 1
//	Primitive 1024th root of unity: 49 (i.e., 49^1024 ≡ 1, 49^512 ≡ -1 mod 12289)
func Falcon512() Params {
	return Params{
		N:           512,
		Q:           12289,
		RootOfUnity: 49,
		MontR:       1 << 16,
		QInv:        qinv(12289, 16),
	}
}

// NewHope256 returns parameters for a NewHope-style ring with n=256.
//
//	R_q = Z_7681[x]/(x^256 + 1)
//	7681 = 15·512 + 1, so 512 | 7680 = q-1
//	Primitive 512th root of unity: 4055
//
// Suitable for KEM (Ring-LWE encryption) at ~128-bit post-quantum security.
// Unlike Kyber's q=3329 which requires an incomplete NTT (512 ∤ 3328),
// q=7681 supports the full negacyclic NTT.
func NewHope256() Params {
	return Params{
		N:           256,
		Q:           7681,
		RootOfUnity: 4055,
		MontR:       1 << 16,
		QInv:        qinv(7681, 16),
	}
}

// HE64 returns parameters for homomorphic evaluation with n=64.
//
//	R_q = Z_10000769[x]/(x^64 + 1)
//	q = 10000769 (24-bit prime, q ≡ 1 mod 128)
//	Primitive 128th root of unity: 6028202
//
// n=64 matches the hamadryad ring dimension. The 24-bit q gives enough
// noise budget for depth-1 multiplication:
//   Fresh noise ≈ 2 * n * eta = 256 (eta=2)
//   Tensor product noise ≈ 4 * n * 256^2 ≈ 16.8M < q/2 ≈ 5M
//
// Wait — that's still too tight. Let's check:
//   With eta=1: noise ≈ 2 * 64 * 1 = 128
//   Product: 4 * 64 * 128^2 = 4,194,304 < 5,000,384 (q/2)  ✓
//
// So eta=1 gives exactly one level of multiplication.
func HE64() Params {
	return Params{
		N:           64,
		Q:           10000769,
		RootOfUnity: 6028202,
		MontR:       1 << 32,
		QInv:        qinv(10000769, 32),
	}
}

// Falcon1024 returns parameters for the Falcon-1024 ring (~256-bit security).
//
//	R_q = Z_12289[x]/(x^1024 + 1)
//	12289-1 = 12288 = 6·2048, so 2048 | 12288
//	Primitive 2048th root of unity: 1945
func Falcon1024() Params {
	return Params{
		N:           1024,
		Q:           12289,
		RootOfUnity: 1945,
		MontR:       1 << 16,
		QInv:        qinv(12289, 16),
	}
}

// qinv computes -Q^{-1} mod 2^k using Hensel lifting.
// Starts from Q^{-1} mod 2 = 1 (Q is odd) and doubles precision each step.
func qinv(q uint32, k int) uint64 {
	// Q is odd, so Q^{-1} mod 2 = 1.
	inv := uint64(1)
	for i := 1; i < k; i++ {
		// Hensel lift: if inv·Q ≡ 1 (mod 2^i), then
		// inv' = inv·(2 - inv·Q) satisfies inv'·Q ≡ 1 (mod 2^{2i}).
		inv = inv * (2 - inv*uint64(q))
	}
	// We want -Q^{-1} mod 2^k.
	mask := (uint64(1) << k) - 1
	return (-inv) & mask
}

// Poly is a polynomial in R_q, stored as N coefficients in [0, Q).
// Can be in either coefficient form or NTT (evaluation) form.
// Operations document which form they expect.
type Poly struct {
	Coeffs []uint32
	params Params
	isNTT  bool
}

// New creates a zero polynomial in coefficient form.
func New(p Params) *Poly {
	return &Poly{
		Coeffs: make([]uint32, p.N),
		params: p,
	}
}

// NewFromCoeffs creates a polynomial from a coefficient slice.
// Coefficients are reduced mod Q.
func NewFromCoeffs(p Params, coeffs []uint32) *Poly {
	poly := New(p)
	n := min(len(coeffs), p.N)
	for i := 0; i < n; i++ {
		poly.Coeffs[i] = coeffs[i] % p.Q
	}
	return poly
}

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

// Params returns the ring parameters.
func (a *Poly) Params() Params { return a.params }

// IsNTT reports whether the polynomial is in NTT form.
func (a *Poly) IsNTT() bool { return a.isNTT }

// addMod computes (a + b) mod q without overflow for uint32 inputs.
func addMod(a, b, q uint32) uint32 {
	sum := a + b
	if sum >= q {
		sum -= q
	}
	return sum
}

// subMod computes (a - b) mod q without overflow.
func subMod(a, b, q uint32) uint32 {
	if a >= b {
		return a - b
	}
	return q - b + a
}

// mulMod computes (a * b) mod q using uint64 intermediate.
func mulMod(a, b, q uint32) uint32 {
	return uint32((uint64(a) * uint64(b)) % uint64(q))
}

// Add sets c = a + b (coefficient-wise mod Q). a and b must be same form.
func Add(a, b *Poly) *Poly {
	c := New(a.params)
	c.isNTT = a.isNTT
	q := a.params.Q
	for i := range a.Coeffs {
		c.Coeffs[i] = addMod(a.Coeffs[i], b.Coeffs[i], q)
	}
	return c
}

// Sub sets c = a - b (coefficient-wise mod Q).
func Sub(a, b *Poly) *Poly {
	c := New(a.params)
	c.isNTT = a.isNTT
	q := a.params.Q
	for i := range a.Coeffs {
		c.Coeffs[i] = subMod(a.Coeffs[i], b.Coeffs[i], q)
	}
	return c
}

// Neg sets c = -a (coefficient-wise mod Q).
func Neg(a *Poly) *Poly {
	c := New(a.params)
	c.isNTT = a.isNTT
	q := a.params.Q
	for i := range a.Coeffs {
		if a.Coeffs[i] != 0 {
			c.Coeffs[i] = q - a.Coeffs[i]
		}
	}
	return c
}

// MulPointwise computes c = a * b pointwise. Both must be in NTT form.
func MulPointwise(a, b *Poly) *Poly {
	c := New(a.params)
	c.isNTT = true
	q := a.params.Q
	for i := range a.Coeffs {
		c.Coeffs[i] = mulMod(a.Coeffs[i], b.Coeffs[i], q)
	}
	return c
}

// MulAdd computes c += a * b pointwise (fused multiply-accumulate in NTT domain).
func MulAdd(c, a, b *Poly) {
	q := a.params.Q
	for i := range a.Coeffs {
		c.Coeffs[i] = addMod(c.Coeffs[i], mulMod(a.Coeffs[i], b.Coeffs[i], q), q)
	}
}

// ScalarMul computes c = s * a (each coefficient multiplied by scalar s mod Q).
func ScalarMul(a *Poly, s uint32) *Poly {
	c := New(a.params)
	c.isNTT = a.isNTT
	q := a.params.Q
	s = s % q
	for i := range a.Coeffs {
		c.Coeffs[i] = mulMod(a.Coeffs[i], s, q)
	}
	return c
}

// Equal reports whether two polynomials have identical coefficients.
func Equal(a, b *Poly) bool {
	if len(a.Coeffs) != len(b.Coeffs) {
		return false
	}
	for i := range a.Coeffs {
		if a.Coeffs[i] != b.Coeffs[i] {
			return false
		}
	}
	return true
}