// Homomorphic evaluation on Ring-LWE ciphertexts (BGV scheme, plaintext mod 2).
package ring

import (
	"crypto/rand"
	"io"

	"crypto/sha3"
)

type HEParams struct {
	KEM      KEMParams
	RelinKey *RelinearizationKey
}

type HECiphertext struct {
	U             *Poly
	V             *Poly
	NoiseEstimate float64
	params        KEMParams
}

type RelinearizationKey struct {
	A []*Poly
	B []*Poly
	L int32
}

func DefaultHEParams() (kp KEMParams) {
	return KEMParams{
		Ring:         HE64(),
		Eta1:         1,
		Eta2:         1,
		SharedKeyLen: 32,
	}
}

func HEKeyGen(kp KEMParams) (pk *KEMPublicKey, sk *KEMSecretKey, rlk *RelinearizationKey) {
	return HEKeyGenFrom(kp, rand.Reader)
}

func HEKeyGenFrom(kp KEMParams, rng io.Reader) (pk *KEMPublicKey, sk *KEMSecretKey, rlk *RelinearizationKey) {
	p := kp.Ring

	a := UniformPolyFrom(p, rng)
	NTT(a)

	s := TernaryPolyFrom(p, rng)
	NTT(s)

	e := CBDPolyFrom(p, kp.Eta1, rng)
	e = ScalarMul(e, 2)
	NTT(e)

	b := MulPointwise(a, s)
	b = Add(b, e)

	z := []byte{:kp.SharedKeyLen}
	io.ReadFull(rng, z)

	pk = &KEMPublicKey{A: a, B: b, P: kp}
	sk = &KEMSecretKey{S: s, PK: pk, Z: z}
	rlk = genRelinKey(sk, rng)
	return pk, sk, rlk
}

func genRelinKey(sk *KEMSecretKey, rng io.Reader) (rlk *RelinearizationKey) {
	p := sk.PK.P.Ring
	q := p.Q

	base := uint32(256)
	levels := int32(0)
	for v := q; v > 0; v /= base {
		levels++
	}

	s2 := MulPointwise(sk.S, sk.S)

	rlkA := []*Poly{:levels}
	rlkB := []*Poly{:levels}

	power := uint32(1)
	for i := int32(0); i < levels; i++ {
		ai := UniformPolyFrom(p, rng)
		NTT(ai)

		ei := CBDPolyFrom(p, sk.PK.P.Eta1, rng)
		ei = ScalarMul(ei, 2)
		NTT(ei)

		bi := MulPointwise(ai, sk.S)
		bi = Add(bi, ei)

		s2Scaled := ScalarMul(s2, power)
		bi = Add(bi, s2Scaled)

		rlkA[i] = ai
		rlkB[i] = bi

		power = mulMod(power, base, q)
	}

	return &RelinearizationKey{A: rlkA, B: rlkB, L: levels}
}

func HEEncrypt(pk *KEMPublicKey, bit int32) (ct *HECiphertext) {
	return HEEncryptFrom(pk, bit, rand.Reader)
}

func HEEncryptFrom(pk *KEMPublicKey, bit int32, rng io.Reader) (ct *HECiphertext) {
	p := pk.P.Ring
	q := p.Q

	coins := []byte{:32}
	io.ReadFull(rng, coins)

	noiseRng := sha3.NewSHAKE256()
	noiseRng.Write(coins)

	r := TernaryPolyFrom(p, noiseRng)
	NTT(r)

	e1 := CBDPolyFrom(p, pk.P.Eta1, noiseRng)
	e2 := CBDPolyFrom(p, pk.P.Eta2, noiseRng)
	e1 = ScalarMul(e1, 2)
	e2 = ScalarMul(e2, 2)

	u := MulPointwise(pk.A, r)
	INTT(u)
	u = Add(u, e1)

	v := MulPointwise(pk.B, r)
	INTT(v)
	v = Add(v, e2)

	if bit != 0 {
		v.Coeffs[0] = addMod(v.Coeffs[0], 1, q)
	}

	return &HECiphertext{
		U:             u,
		V:             v,
		NoiseEstimate: float64(pk.P.Eta1+pk.P.Eta2) * 4,
		params:        pk.P,
	}
}

func HEDecrypt(sk *KEMSecretKey, ct *HECiphertext) (result int32) {
	uNTT := ct.U.Clone()
	NTT(uNTT)
	su := MulPointwise(sk.S, uNTT)
	INTT(su)

	noisy := Sub(ct.V, su)

	q := ct.params.Ring.Q
	c := noisy.Coeffs[0]

	if c > q/2 {
		c = q - c
	}
	return int32(c % 2)
}

func HEAdd(a, b *HECiphertext) (c *HECiphertext) {
	return &HECiphertext{
		U:             Add(a.U, b.U),
		V:             Add(a.V, b.V),
		NoiseEstimate: a.NoiseEstimate + b.NoiseEstimate,
		params:        a.params,
	}
}

func HESub(a, b *HECiphertext) (c *HECiphertext) {
	return &HECiphertext{
		U:             Sub(a.U, b.U),
		V:             Sub(a.V, b.V),
		NoiseEstimate: a.NoiseEstimate + b.NoiseEstimate,
		params:        a.params,
	}
}

func HEMul(a, b *HECiphertext, rlk *RelinearizationKey) (result *HECiphertext) {
	p := a.params.Ring

	c0 := Mul(a.V, b.V)
	c1a := Mul(a.V, b.U)
	c1b := Mul(a.U, b.V)
	c1 := Add(c1a, c1b)
	c2 := Mul(a.U, b.U)

	base := uint32(256)
	relinU := New(p)
	relinV := New(p)

	if c2.isNTT {
		INTT(c2)
	}

	for level := int32(0); level < rlk.L; level++ {
		digit := New(p)
		for j := int32(0); j < int32(len(digit.Coeffs)); j++ {
			d := c2.Coeffs[j]
			for k := int32(0); k < level; k++ {
				d /= base
			}
			digit.Coeffs[j] = d % base
		}

		NTT(digit)

		uPart := MulPointwise(digit, rlk.A[level])
		vPart := MulPointwise(digit, rlk.B[level])
		INTT(uPart)
		INTT(vPart)

		relinU = Add(relinU, uPart)
		relinV = Add(relinV, vPart)
	}

	return &HECiphertext{
		U:             Add(c1, relinU),
		V:             Add(c0, relinV),
		NoiseEstimate: a.NoiseEstimate * b.NoiseEstimate * 4,
		params:        a.params,
	}
}

func HENot(a *HECiphertext) (result *HECiphertext) {
	p := a.params.Ring

	newV := a.V.Clone()
	newV.Coeffs[0] = addMod(newV.Coeffs[0], 1, p.Q)

	return &HECiphertext{
		U:             a.U.Clone(),
		V:             newV,
		NoiseEstimate: a.NoiseEstimate,
		params:        a.params,
	}
}

func HEXOR(a, b *HECiphertext) (c *HECiphertext) {
	return HEAdd(a, b)
}

func HEAND(a, b *HECiphertext, rlk *RelinearizationKey) (c *HECiphertext) {
	return HEMul(a, b, rlk)
}

func Rerandomize(pk *KEMPublicKey, ct *HECiphertext) (result *HECiphertext) {
	return RerandomizeFrom(pk, ct, rand.Reader)
}

func RerandomizeFrom(pk *KEMPublicKey, ct *HECiphertext, rng io.Reader) (result *HECiphertext) {
	zero := HEEncryptFrom(pk, 0, rng)
	return HEAdd(ct, zero)
}