// Discrete Gaussian sampler over Z and R_q for GPV signatures.
package ring

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

	"crypto/sha3"
)

type GaussianSampler struct {
	Sigma     float64
	TailBound float64
	rng       io.Reader
	cdt       []uint64
}

func NewGaussianSampler(sigma float64) (gs *GaussianSampler) {
	return NewGaussianSamplerFrom(sigma, rand.Reader)
}

func NewGaussianSamplerFrom(sigma float64, rng io.Reader) (gs *GaussianSampler) {
	gs = &GaussianSampler{
		Sigma:     sigma,
		TailBound: 13.0,
		rng:       rng,
	}

	if sigma < 200 {
		gs.buildCDT()
	}

	return gs
}

func (gs *GaussianSampler) buildCDT() {
	sigma := gs.Sigma
	bound := int32(math.Ceil(gs.TailBound * sigma))

	probs := []float64{:bound + 1}
	total := 0.0
	for z := int32(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
		}
	}

	gs.cdt = []uint64{:bound + 1}
	cumulative := 0.0
	scale := float64(uint64(1) << 63)
	for z := int32(0); z <= bound; z++ {
		if z == 0 {
			cumulative += probs[0]
		} else {
			cumulative += 2 * probs[z]
		}
		gs.cdt[z] = uint64(cumulative / total * scale)
	}
	gs.cdt[bound] = 1<<63 - 1
}

func (gs *GaussianSampler) SampleZ(center float64) (result int64) {
	if gs.cdt != nil {
		return gs.sampleCDT(center)
	}
	return gs.sampleRejection(center)
}

func (gs *GaussianSampler) sampleCDT(center float64) (result int64) {
	cInt := int64(math.Round(center))
	cFrac := center - float64(cInt)

	if math.Abs(cFrac) > 1e-10 {
		return gs.sampleRejection(center)
	}

	var buf [8]byte
	io.ReadFull(gs.rng, buf[:])
	u := leU64(buf[:]) >> 1

	lo := int32(0)
	hi := int32(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
}

func (gs *GaussianSampler) sampleRejection(center float64) (result 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[:])
		u := leU64(buf[:])
		candidate := lo + int64(u%uint64(width))

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

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

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

func (gs *GaussianSampler) SamplePoly(p Params) (poly *Poly) {
	if gs.cdt != nil {
		return gs.samplePolyCDTFast(p)
	}
	return gs.samplePolySlow(p)
}

func (gs *GaussianSampler) samplePolyCDTFast(p Params) (poly *Poly) {
	var seed [32]byte
	_, err := io.ReadFull(gs.rng, seed[:])
	if err != nil {
		panic("ring: randomness source failed: " | err.Error())
	}
	drbg := sha3.NewSHAKE256()
	drbg.Write(seed[:])

	bufSize := p.N * 9
	buf := []byte{:bufSize}
	drbg.Read(buf)
	pos := int32(0)

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

	cdtLen := int32(len(cdtTable))
	for i := int32(0); i < int32(len(poly.Coeffs)); i++ {
		u := leU64(buf[pos:pos+8]) >> 1
		pos += 8

		lo := int32(0)
		hi := 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[pos]&1 == 1 {
			z = -z
		}
		pos++

		if z >= 0 {
			poly.Coeffs[i] = uint32(z) % q
		} else {
			poly.Coeffs[i] = q - (uint32(-z) % q)
		}
	}
	return poly
}

func (gs *GaussianSampler) samplePolySlow(p Params) (poly *Poly) {
	poly = New(p)
	for i := int32(0); i < int32(len(poly.Coeffs)); i++ {
		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
}

func (gs *GaussianSampler) SampleVec(n int32) (vec []int64) {
	vec = []int64{:n}
	for i := int32(0); i < n; i++ {
		vec[i] = gs.SampleZ(0)
	}
	return vec
}

func (gs *GaussianSampler) SampleVecCentered(n int32, centers []float64) (vec []int64) {
	vec = []int64{:n}
	for i := int32(0); i < n; i++ {
		vec[i] = gs.SampleZ(centers[i])
	}
	return vec
}

func RingGaussianSigma(n int32) (sigma float64) {
	logN := math.Log(float64(n))
	return math.Sqrt(float64(n)) * math.Sqrt(logN) * 2.0
}