// Discrete Gaussian sampler over Z and R_q for GPV signatures.
//
// The discrete Gaussian distribution D_{Z, σ, c} over the integers assigns
// probability proportional to exp(-π(x-c)²/σ²) to each x ∈ Z.
//
// For lattice-based cryptography, we need:
//   - D_{Z, σ}:    1-dimensional Gaussian over integers (base case)
//   - D_{Z^n, σ}:  n-dimensional by sampling each coordinate independently
//   - D_{Λ, σ, c}: Gaussian over lattice coset (used in GPV signing)
//
// This implements the Peikert convolution technique and the
// rejection-sampling method of Ducas, Durmus, Lepoint, Lyubashevsky (2013).
//
// The parameter σ (sigma) controls the width of the distribution:
//   - Larger σ → more spread → more secure but larger signatures
//   - σ must exceed the smoothing parameter η_ε(Λ) for statistical closeness
//   - For Falcon-512: σ ≈ 165 (at ~128-bit security)
//
// Security: the sampler must produce outputs statistically close to a true
// discrete Gaussian. Any bias leaks information about the trapdoor.

package ring

import (
	"crypto/rand"
	"encoding/binary"
	"io"
	"math"

	"golang.org/x/crypto/sha3"
)

// GaussianSampler samples from the discrete Gaussian D_{Z, σ, c}.
type GaussianSampler struct {
	// Sigma is the Gaussian parameter (standard deviation ≈ σ/√(2π)).
	Sigma float64

	// TailBound is the cutoff: |x - c| ≤ TailBound * σ.
	// Beyond this, probability is negligible (< 2^{-128} for TailBound=13).
	TailBound float64

	// rng is the randomness source.
	rng io.Reader

	// Precomputed table for the CDT (Cumulative Distribution Table) method.
	// Only used for small σ (σ < 10). For larger σ, use rejection sampling.
	cdt []uint64
}

// NewGaussianSampler creates a sampler for D_{Z, σ, 0}.
func NewGaussianSampler(sigma float64) *GaussianSampler {
	return NewGaussianSamplerFrom(sigma, rand.Reader)
}

// NewGaussianSamplerFrom creates a sampler with the given randomness source.
func NewGaussianSamplerFrom(sigma float64, rng io.Reader) *GaussianSampler {
	gs := &GaussianSampler{
		Sigma:     sigma,
		TailBound: 13.0, // 2^{-128} security
		rng:       rng,
	}

	// Build CDT for σ up to ~200 (Falcon-512 range).
	// The table has ceil(13·σ) entries — 424 for σ=32.6, 2145 for σ=165.
	// Binary search makes lookup O(log n) ≈ 9-11 comparisons regardless.
	// Only fall back to rejection for σ > 200 where the table exceeds ~2600 entries.
	if sigma < 200 {
		gs.buildCDT()
	}

	return gs
}

// buildCDT constructs the cumulative distribution table.
// For each z ≥ 0, CDF[z] = Σ_{x=0}^{z} ρ_σ(x), scaled to [0, 2^63).
func (gs *GaussianSampler) buildCDT() {
	sigma := gs.Sigma
	bound := int(math.Ceil(gs.TailBound * sigma))

	// Compute unnormalized probabilities for z = 0, 1, ..., bound.
	probs := make([]float64, bound+1)
	total := 0.0
	for z := 0; z <= bound; z++ {
		p := math.Exp(-math.Pi * float64(z) * float64(z) / (sigma * sigma))
		probs[z] = p
		if z == 0 {
			total += p
		} else {
			total += 2 * p // both +z and -z
		}
	}

	// Build CDF scaled to 2^63.
	gs.cdt = make([]uint64, bound+1)
	cumulative := 0.0
	scale := float64(uint64(1) << 63)
	for z := 0; z <= bound; z++ {
		// P(|X| ≤ z) = sum of probs for 0..z, where z>0 counts twice
		if z == 0 {
			cumulative += probs[0]
		} else {
			cumulative += 2 * probs[z]
		}
		gs.cdt[z] = uint64(cumulative / total * scale)
	}
	// Ensure last entry is max.
	gs.cdt[bound] = 1<<63 - 1
}

// SampleZ samples from D_{Z, σ, c} (discrete Gaussian centered at c).
func (gs *GaussianSampler) SampleZ(center float64) int64 {
	if gs.cdt != nil {
		return gs.sampleCDT(center)
	}
	return gs.sampleRejection(center)
}

// sampleCDT uses the cumulative distribution table for small σ.
// Samples |z| from the table, then randomly assigns sign.
func (gs *GaussianSampler) sampleCDT(center float64) int64 {
	// Sample from D_{Z, σ, 0} then shift by round(center).
	cInt := int64(math.Round(center))
	cFrac := center - float64(cInt)

	// If fractional part is significant, fall back to rejection.
	if math.Abs(cFrac) > 1e-10 {
		return gs.sampleRejection(center)
	}

	// Sample |z| using CDT with binary search.
	var buf [8]byte
	io.ReadFull(gs.rng, buf[:])
	u := binary.LittleEndian.Uint64(buf[:]) >> 1 // 63-bit uniform

	// Binary search: find smallest z where cdt[z] > u.
	lo, hi := 0, len(gs.cdt)-1
	for lo < hi {
		mid := (lo + hi) / 2
		if gs.cdt[mid] <= u {
			lo = mid + 1
		} else {
			hi = mid
		}
	}
	z := int64(lo)

	// Random sign for z > 0.
	if z > 0 {
		io.ReadFull(gs.rng, buf[:1])
		if buf[0]&1 == 1 {
			z = -z
		}
	}

	return z + cInt
}

// sampleRejection uses rejection sampling for arbitrary σ and center.
// Samples a candidate from a uniform distribution over the tail bound,
// then accepts with probability proportional to ρ_σ(x - c).
func (gs *GaussianSampler) sampleRejection(center float64) int64 {
	sigma := gs.Sigma
	bound := int64(math.Ceil(gs.TailBound * sigma))
	lo := int64(math.Floor(center)) - bound
	hi := int64(math.Ceil(center)) + bound
	width := hi - lo + 1

	var buf [8]byte
	piOverSigma2 := math.Pi / (sigma * sigma)

	for {
		io.ReadFull(gs.rng, buf[:])
		// Sample candidate uniformly from [lo, hi].
		u := binary.LittleEndian.Uint64(buf[:])
		candidate := lo + int64(u%uint64(width))

		// Accept with probability exp(-π(candidate - center)² / σ²).
		diff := float64(candidate) - center
		logProb := -piOverSigma2 * diff * diff

		// Sample a uniform [0, 1) and check.
		io.ReadFull(gs.rng, buf[:])
		uFloat := float64(binary.LittleEndian.Uint64(buf[:])>>11) / float64(uint64(1)<<53)

		if math.Log(uFloat) < logProb {
			return candidate
		}
	}
}

// SamplePoly samples a polynomial with each coefficient drawn from D_{Z, σ, 0}.
//
// Uses a SHAKE256 DRBG seeded from one 32-byte crypto/rand read, with all
// randomness squeezed into a single buffer upfront. This reduces kernel
// crossings from 128+ to 1 and eliminates per-sample interface dispatch.
// SHAKE256 output is computationally indistinguishable from random
// (it's the PRF used in ML-KEM/Kyber).
func (gs *GaussianSampler) SamplePoly(p Params) *Poly {
	if gs.cdt != nil {
		return gs.samplePolyCDTFast(p)
	}
	return gs.samplePolySlow(p)
}

// samplePolyCDTFast is the optimized CDT path with bulk randomness.
func (gs *GaussianSampler) samplePolyCDTFast(p Params) *Poly {
	// Seed a SHAKE256 XOF from one crypto/rand read.
	var seed [32]byte
	if _, err := io.ReadFull(gs.rng, seed[:]); err != nil {
		panic("ring: randomness source failed: " + err.Error())
	}
	drbg := sha3.NewShake256()
	drbg.Write(seed[:])

	// Pre-squeeze all randomness: 9 bytes per coefficient (8 for CDT + 1 for sign).
	// Over-allocate slightly for rejection fallback (rare with CDT).
	bufSize := p.N * 9
	buf := make([]byte, bufSize)
	drbg.Read(buf)
	pos := 0

	poly := New(p)
	cdtTable := gs.cdt
	q := p.Q

	cdtLen := len(cdtTable)
	for i := range poly.Coeffs {
		// Inline CDT sampling with binary search.
		u := binary.LittleEndian.Uint64(buf[pos:pos+8]) >> 1 // 63-bit uniform
		pos += 8

		// Binary search: find smallest z where cdtTable[z] > u.
		lo, hi := 0, cdtLen-1
		for lo < hi {
			mid := (lo + hi) / 2
			if cdtTable[mid] <= u {
				lo = mid + 1
			} else {
				hi = mid
			}
		}
		z := int64(lo)

		// Random sign for z > 0.
		if z > 0 && buf[pos]&1 == 1 {
			z = -z
		}
		pos++

		// Convert to centered representation mod Q.
		if z >= 0 {
			poly.Coeffs[i] = uint32(z) % q
		} else {
			poly.Coeffs[i] = q - (uint32(-z) % q)
		}
	}
	return poly
}

// samplePolySlow is the fallback for large σ (rejection sampling).
func (gs *GaussianSampler) samplePolySlow(p Params) *Poly {
	poly := New(p)
	for i := range poly.Coeffs {
		z := gs.SampleZ(0)
		if z >= 0 {
			poly.Coeffs[i] = uint32(z) % p.Q
		} else {
			poly.Coeffs[i] = p.Q - (uint32(-z) % p.Q)
		}
	}
	return poly
}

// SampleVec samples a vector of n independent Gaussian integers.
func (gs *GaussianSampler) SampleVec(n int) []int64 {
	vec := make([]int64, n)
	for i := range vec {
		vec[i] = gs.SampleZ(0)
	}
	return vec
}

// SampleVecCentered samples a vector of n Gaussian integers centered at c.
func (gs *GaussianSampler) SampleVecCentered(n int, centers []float64) []int64 {
	vec := make([]int64, n)
	for i := range vec {
		vec[i] = gs.SampleZ(centers[i])
	}
	return vec
}

// RingGaussianSigma returns the recommended σ for a given ring dimension
// to achieve security level λ bits.
//
// For GPV signatures over R_q with n coefficients:
//   σ ≈ s1(R) · ω(√log n)
// where s1(R) is the largest singular value of the trapdoor basis.
//
// Practical values:
//   n=512:  σ ≈ 165 (Falcon-512 target)
//   n=1024: σ ≈ 168 (Falcon-1024 target)
//   n=64:   σ ≈ 50  (for testing)
func RingGaussianSigma(n int) float64 {
	// σ = √(n) * √(log(n)) * constant
	// The constant is tuned so that σ/q gives appropriate noise ratio.
	logN := math.Log(float64(n))
	return math.Sqrt(float64(n)) * math.Sqrt(logN) * 2.0
}