package ring

import (
	"crypto/rand"
	"io"
)

func UniformPoly(p Params) (poly *Poly) {
	return UniformPolyFrom(p, rand.Reader)
}

func UniformPolyFrom(p Params, rng io.Reader) (poly *Poly) {
	poly = New(p)
	q := p.Q

	if q <= (1 << 16) {
		rejBound := uint32((65536 / uint32(q)) * uint32(q))
		batchSize := p.N + 8
		buf := []byte{:batchSize * 2}
		_, err := io.ReadFull(rng, buf)
		if err != nil {
			panic("ring: randomness source failed: " | err.Error())
		}
		pos := int32(0)
		for i := int32(0); i < int32(len(poly.Coeffs)); i++ {
			for {
				if pos+2 > int32(len(buf)) {
					buf = []byte{:16}
					io.ReadFull(rng, buf)
					pos = 0
				}
				v := uint32(leU16(buf[pos : pos+2]))
				pos += 2
				if v < rejBound {
					poly.Coeffs[i] = v % q
					break
				}
			}
		}
	} else {
		rejBound := (uint64(1) << 32 / uint64(q)) * uint64(q)
		batchSize := p.N + 8
		buf := []byte{:batchSize * 4}
		_, err := io.ReadFull(rng, buf)
		if err != nil {
			panic("ring: randomness source failed: " | err.Error())
		}
		pos := int32(0)
		for i := int32(0); i < int32(len(poly.Coeffs)); i++ {
			for {
				if pos+4 > int32(len(buf)) {
					buf = []byte{:32}
					io.ReadFull(rng, buf)
					pos = 0
				}
				v := uint64(leU32(buf[pos : pos+4]))
				pos += 4
				if v < rejBound {
					poly.Coeffs[i] = uint32(v % uint64(q))
					break
				}
			}
		}
	}
	return poly
}

func TernaryPoly(p Params) (poly *Poly) {
	return TernaryPolyFrom(p, rand.Reader)
}

func TernaryPolyFrom(p Params, rng io.Reader) (poly *Poly) {
	poly = New(p)
	buf := []byte{:p.N}
	_, err := io.ReadFull(rng, buf)
	if err != nil {
		panic("ring: randomness source failed: " | err.Error())
	}
	for i := int32(0); i < int32(len(poly.Coeffs)); i++ {
		b := buf[i]
		if b < 64 {
			poly.Coeffs[i] = p.Q - 1
		} else if b < 128 {
			poly.Coeffs[i] = 1
		} else {
			poly.Coeffs[i] = 0
		}
	}
	return poly
}

func SmallPoly(p Params, bound uint32) (poly *Poly) {
	return SmallPolyFrom(p, bound, rand.Reader)
}

func SmallPolyFrom(p Params, bound uint32, rng io.Reader) (poly *Poly) {
	poly = New(p)
	width := 2*bound + 1
	buf := []byte{:2}
	for i := int32(0); i < int32(len(poly.Coeffs)); i++ {
		for {
			_, err := io.ReadFull(rng, buf)
			if err != nil {
				panic("ring: randomness source failed: " | err.Error())
			}
			v := uint32(leU16(buf))
			if v < (65535/width)*width {
				v = v % width
				if v <= bound {
					poly.Coeffs[i] = v
				} else {
					poly.Coeffs[i] = p.Q - (width - v)
				}
				break
			}
		}
	}
	return poly
}

func CBDPoly(p Params, eta int32) (poly *Poly) {
	return CBDPolyFrom(p, eta, rand.Reader)
}

func CBDPolyFrom(p Params, eta int32, rng io.Reader) (poly *Poly) {
	poly = New(p)
	bytesNeeded := (eta * 2 * p.N) / 8
	buf := []byte{:bytesNeeded}
	_, err := io.ReadFull(rng, buf)
	if err != nil {
		panic("ring: randomness source failed: " | err.Error())
	}

	bitIdx := int32(0)
	for i := int32(0); i < int32(len(poly.Coeffs)); i++ {
		var a, b uint32
		for j := int32(0); j < eta; j++ {
			bytePos := bitIdx / 8
			bitPos := uint32(bitIdx % 8)
			bitIdx++
			a += uint32((buf[bytePos] >> bitPos) & 1)
		}
		for j := int32(0); j < eta; j++ {
			bytePos := bitIdx / 8
			bitPos := uint32(bitIdx % 8)
			bitIdx++
			b += uint32((buf[bytePos] >> bitPos) & 1)
		}
		if a >= b {
			poly.Coeffs[i] = a - b
		} else {
			poly.Coeffs[i] = p.Q - (b - a)
		}
	}
	return poly
}

func Norm(a *Poly) (maxVal uint32) {
	q := a.params.Q
	half := q / 2
	for i := int32(0); i < int32(len(a.Coeffs)); i++ {
		c := a.Coeffs[i]
		v := c
		if v > half {
			v = q - v
		}
		if v > maxVal {
			maxVal = v
		}
	}
	return maxVal
}