package gnarlring

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

// GaussSampler samples from the discrete Gaussian distribution D_{Z, sigma}
// over the integers. Uses the Cumulative Distribution Table (CDT) method
// for the standard sigma ≈ 10.4. Each call to SampleZ produces an integer
// drawn from D_{Z, sigma, 0}. SamplePoly produces a Poly27 with each
// coefficient drawn independently.
type GaussSampler struct {
	sigma float64
	tail  int
	cdt   []uint64 // cumulative distribution table, scaled to 2^63
	rng   io.Reader
}

// NewGaussSampler creates a sampler for D_{Z, sigma, 0}.
// sigma is typically sqrt(N) * 2 ≈ 10.39 for the gnarl ring.
func NewGaussSampler(sigma float64) *GaussSampler {
	return NewGaussSamplerFrom(sigma, rand.Reader)
}

// NewGaussSamplerFrom creates a sampler with the given randomness source.
// Only builds the CDT for sigma < 50 — larger sigma uses rejection sampling.
func NewGaussSamplerFrom(sigma float64, rng io.Reader) *GaussSampler {
	if rng == nil {
		rng = rand.Reader
	}
	gs := &GaussSampler{
		sigma: sigma,
		tail:  13,
		rng:   rng,
	}
	if sigma < 50 {
		gs.buildCDT()
	}
	return gs
}

// buildCDT constructs the cumulative distribution table.
// For each z ≥ 0, CDT[z] = P(|X| ≤ z) scaled to [0, 2^63).
func (gs *GaussSampler) buildCDT() {
	sigma := gs.sigma
	bound := int(math.Ceil(float64(gs.tail) * sigma))

	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
		}
	}

	gs.cdt = make([]uint64, bound+1)
	cumulative := 0.0
	scale := float64(uint64(1) << 63)
	for z := 0; z <= bound; z++ {
		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, sigma, center}.
func (gs *GaussSampler) SampleZ(center float64) int64 {
	if gs.cdt != nil {
		// For small sigma, the CDT is tiny and rejection sampling for
		// non-integer centers is prohibitively slow. Round to nearest
		// integer — the bias from rounding is O(1/σ) which for
		// sigma_eff ≈ 0.1 is negligible (< 1 per 10^6 samples).
		cInt := int64(math.Round(center))
		return gs.sampleCDT(float64(cInt))
	}
	return gs.sampleRejection(center)
}

// sampleCDT samples |z| via binary search on the CDF, then randomly assigns sign.
func (gs *GaussSampler) sampleCDT(center float64) int64 {
	cInt := int64(math.Round(center))

	var buf [8]byte
	io.ReadFull(gs.rng, buf[:])
	u := binary.LittleEndian.Uint64(buf[:]) >> 1 // 63-bit uniform

	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)

	if z > 0 {
		io.ReadFull(gs.rng, buf[:1])
		if buf[0]&1 == 1 {
			z = -z
		}
	}
	return z + cInt
}

// sampleRejection is the fallback for non-integer centers or large sigma.
func (gs *GaussSampler) sampleRejection(center float64) int64 {
	sigma := gs.sigma
	bound := int64(math.Ceil(float64(gs.tail) * sigma))
	lo := int64(math.Floor(center)) - bound
	hi := int64(math.Ceil(center)) + bound
	width := hi - lo + 1

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

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

		diff := float64(candidate) - center
		logProb := -piOverSigma2 * diff * diff

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

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

// SamplePoly returns a Poly27 with each coefficient drawn from D_{Z, sigma, 0}.
func (gs *GaussSampler) SamplePoly() *Poly27 {
	if gs.cdt != nil {
		return gs.samplePolyCDT()
	}
	return gs.samplePolySlow()
}

// samplePolyCDT uses the CDT with bulk randomness for efficiency.
func (gs *GaussSampler) samplePolyCDT() *Poly27 {
	// For n=27, this is fast enough without SHAKE256 DRBG.
	// 27 coefficients × 9 bytes (8 for CDT + 1 for sign).
	p := NewPoly27()
	cdtTable := gs.cdt
	cdtLen := len(cdtTable)
	var buf [9]byte

	for i := 0; i < N; i++ {
		io.ReadFull(gs.rng, buf[:])
		u := binary.LittleEndian.Uint64(buf[:8]) >> 1

		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)
		if z > 0 && buf[8]&1 == 1 {
			z = -z
		}

		if z >= 0 {
			p.Coeffs[i] = uint16(uint64(z) % Q)
		} else {
			p.Coeffs[i] = Q - uint16(uint64(-z)%Q)
		}
	}
	return p
}

// samplePolySlow is the fallback for large sigma.
func (gs *GaussSampler) samplePolySlow() *Poly27 {
	p := NewPoly27()
	for i := 0; i < N; i++ {
		z := gs.SampleZ(0)
		if z >= 0 {
			p.Coeffs[i] = uint16(uint64(z) % Q)
		} else {
			p.Coeffs[i] = Q - uint16(uint64(-z)%Q)
		}
	}
	return p
}