/*
 * An implementation of Groth and Kohlweiss, "One-out-of-Many Proofs: Or
 * How to Leak a Secret and Spend a Coin", Eurocrypt 2015
 * https://eprint.iacr.org/2014/764.pdf
 */

package main

import (
	"fmt"
	"go.dedis.ch/kyber"
	"go.dedis.ch/kyber/util/random"
)

type GroupParams struct {
	/* A group of prime order */
	group kyber.Group
	/* Three generators of the group; the relative discrete logs
	 * must be unknown. Commitments to a value x with randomness r
	 * are x*A + r*B.  Y is used to pad the list of commitments to a
	 * power of 2. */
	A, B, Y kyber.Point
}

type PubState struct {
	cl, ca, cb, cd []kyber.Point
}

type PrivState struct {
	r, a, s, t, rho []kyber.Scalar
	ell uint32
	privkey kyber.Scalar
}

type Proof struct {
	f, za, zb []kyber.Scalar
	zd kyber.Scalar
}

// Multiply a polynomial expressed as a slice of coefficients by the
// scalar a
func polymul(group kyber.Group, poly []kyber.Scalar, a kyber.Scalar) {
	numcoeffs := len(poly)
	for i := 0 ; i < numcoeffs ; i++ {
		poly[i] = group.Scalar().Mul(poly[i], a)
	}
}

// Multiply a polynomial expressed as a slice of coefficients by
// the linear polynomial x+a
func polymul_xplus(group kyber.Group, poly []kyber.Scalar, a kyber.Scalar) {
	numcoeffs := len(poly)
	if !poly[numcoeffs-1].Equal(group.Scalar().Zero()) {
		panic("Degree too large in multiplication")
	}
	for i := numcoeffs-1 ; i > 0 ; i-- {
		poly[i] = group.Scalar().Add(
			group.Scalar().Mul(poly[i], a), poly[i-1])
	}
	poly[0] = group.Scalar().Mul(poly[0], a)
}

/* Sigma protocol for one out of N commitments containing 0
 * Section 3 of https://eprint.iacr.org/2014/764.pdf
 */
func ProofStep1(params GroupParams, c []kyber.Point, ell uint32, privkey kyber.Scalar) (PubState, PrivState) {
	N := uint32(len(c))
	rand := random.New()
	group := params.group

	// Find n := ceil(log_2(N))
	n := uint32(1)
	two_n := uint32(2) // 2^n
	for ; two_n < N ; {
		n += 1
		two_n *= 2
	}
	fmt.Printf("n = %d\n", n);

	// We will keep the notation of the paper:
	// r, a, s, t, cl, ca, cb are 1-based,
	// c, rho, cd are 0-based.  The 1-based ones, we
	// just allocate an extra array element, and never
	// use element 0.
	var pub PubState
	var priv PrivState
	priv.r = make([]kyber.Scalar, n+1)
	priv.a = make([]kyber.Scalar, n+1)
	priv.s = make([]kyber.Scalar, n+1)
	priv.t = make([]kyber.Scalar, n+1)
	priv.rho = make([]kyber.Scalar, n)
	priv.ell = ell
	priv.privkey = privkey.Clone()
	pub.cl = make([]kyber.Point, n+1)
	pub.ca = make([]kyber.Point, n+1)
	pub.cb = make([]kyber.Point, n+1)
	pub.cd = make([]kyber.Point, n)

	var j, k, mask uint32
	// mask = 2^k
	j = 1
	k = 0
	mask = 1
	for ; k < n ; {
		priv.r[j] = group.Scalar().Pick(rand)
		priv.a[j] = group.Scalar().Pick(rand)
		priv.s[j] = group.Scalar().Pick(rand)
		priv.t[j] = group.Scalar().Pick(rand)
		priv.rho[k] = group.Scalar().Pick(rand)

		pub.cl[j] = group.Point().Mul(priv.r[j], params.B)
		if (ell & mask) != 0 {
			pub.cl[j] = group.Point().Add(pub.cl[j], params.A)
		}
		pub.ca[j] = group.Point().Add(
				group.Point().Mul(priv.a[j], params.A),
				group.Point().Mul(priv.s[j], params.B))
		pub.cb[j] = group.Point().Mul(priv.t[j], params.B)
		if (ell & mask) != 0 {
			pub.cb[j] = group.Point().Add(pub.cb[j],
					group.Point().Mul(priv.a[j], params.A))
		}

		j++
		k++
		mask *= 2
	}
	for k = 0 ; k < n ; k++ {
		pub.cd[k] = group.Point().Mul(priv.rho[k], params.B)
	}

	for i := uint32(0); i < two_n; i++ {
		// Compute the coefficients of p_i
		p_i := make([]kyber.Scalar, n+1)
		p_i[0] = group.Scalar().One()
		for t := uint32(1); t <= n; t++ {
			p_i[t] = group.Scalar().Zero()
		}

		j = 1
		// jmask = 2^(j-1)
		jmask := uint32(1)
		for ; j <= n ; {
			if (i & jmask) != 0 {
				if (ell & jmask) != 0 {
					// Multiply p_i by x + a[j]
					polymul_xplus(group, p_i, priv.a[j])
				} else {
					// Multiply p_i by a[j]
					polymul(group, p_i, priv.a[j])
				}
			} else {
				negaj := group.Scalar().Neg(priv.a[j])
				if (ell & jmask) != 0 {
					// Multiply p_i by -a[j]
					polymul(group, p_i, negaj)
				} else {
					// Multiply p_i by x - a[j]
					polymul_xplus(group, p_i, negaj)
				}
			}

			j++
			jmask *= 2
		}

		if i == ell && !p_i[n].Equal(group.Scalar().One()) {
			panic("Leading coeff should be 1 but was not")
		}
		if i != ell && !p_i[n].Equal(group.Scalar().Zero()) {
			panic("Leading coeff should be 0 but was not")
		}
		for k = 0 ; k < n ; k++ {
			if i < N {
				pub.cd[k] = group.Point().Add(pub.cd[k],
						group.Point().Mul(p_i[k], c[i]))
			} else {
				pub.cd[k] = group.Point().Add(pub.cd[k],
						group.Point().Mul(p_i[k], params.Y))
			}
		}
	}

	return pub, priv
}

func GenChallenge(params GroupParams, pub PubState) kyber.Scalar {
	// In the interactive version, just pick a random challenge.
	// In the noninteractive version, this would be a hash of pub
	// and a message.
	rand := random.New()
	return params.group.Scalar().Pick(rand)
}

func ProofStep2(params GroupParams, priv PrivState, x kyber.Scalar) Proof {
	var proof Proof
	n := uint32(len(priv.rho))
	group := params.group
	proof.f = make([]kyber.Scalar, n+1)
	proof.za = make([]kyber.Scalar, n+1)
	proof.zb = make([]kyber.Scalar, n+1)

	var j, mask uint32
	// mask = 2^(j-1)
	j = 1
	mask = 1
	for ; j <= n ; {
		if (priv.ell & mask) != 0 {
			proof.f[j] = group.Scalar().Add(x, priv.a[j])
		} else {
			proof.f[j] = priv.a[j].Clone()
		}
		proof.za[j] = group.Scalar().Add(
			group.Scalar().Mul(x, priv.r[j]), priv.s[j])
		proof.zb[j] = group.Scalar().Add(
			group.Scalar().Mul(
				group.Scalar().Sub(x, proof.f[j]),
				priv.r[j]),
			priv.t[j])

		j++
		mask *= 2
	}

	proof.zd = group.Scalar().Zero()
	k := uint32(0)
	xk := group.Scalar().One() // x^k
	for ; k < n ; {
		proof.zd = group.Scalar().Sub(proof.zd,
			group.Scalar().Mul(priv.rho[k], xk))
		k++
		xk = group.Scalar().Mul(xk, x)
	}
	// At this point, xk = x^n
	proof.zd = group.Scalar().Add(proof.zd,
		group.Scalar().Mul(priv.privkey, xk))

	return proof
}

// We assume that the checks that the Points are in fact in the right
// group and the Scalars are in the right range were done at the time
// the PubState and Proof structures were populated from data received
// over the network.  This function checks the relations between the
// elements.
func Verify(params GroupParams, c []kyber.Point, pub PubState, x kyber.Scalar, proof Proof) bool {
	N := uint32(len(c))
	group := params.group

	// Find n := ceil(log_2(N))
	n := uint32(1)
	two_n := uint32(2) // 2^n
	for ; two_n < N ; {
		n += 1
		two_n *= 2
	}

	// Check that the lengths of the arrays are correct
	if len(pub.cl) != int(n+1) || len(pub.ca) != int(n+1) ||
			len(pub.cb) != int(n+1) ||
			len(pub.cd) != int(n) ||
			len(proof.f) != int(n+1) ||
			len(proof.za) != int(n+1) ||
			len(proof.zb) != int(n+1) {
		fmt.Printf("Inputs not of the correct length\n")
		return false
	}

	for j := uint32(1) ; j <= n ; j++ {
		lhs1 := group.Point().Add(group.Point().Mul(x, pub.cl[j]),
			pub.ca[j])
		rhs1 := group.Point().Add(
			group.Point().Mul(proof.f[j], params.A),
			group.Point().Mul(proof.za[j], params.B))
		if !lhs1.Equal(rhs1) {
			fmt.Printf("Failed equality check 1\n")
			return false
		}
		lhs2 := group.Point().Add(
			group.Point().Mul(
				group.Scalar().Sub(x, proof.f[j]),
				pub.cl[j]),
			pub.cb[j])
		rhs2 := group.Point().Mul(proof.zb[j], params.B)
		if !lhs2.Equal(rhs2) {
			fmt.Printf("Failed equality check 2\n")
			return false
		}

	}

	lhs := group.Point().Null()
	for i := uint32(0) ; i < two_n ; i++ {
		fprod := group.Scalar().One()
		j := uint32(1)
		mask := uint32(1) // mask = 2^(j-1)
		for ; j <= n ; {
			if (i & mask) != 0 {
				fprod = group.Scalar().Mul(fprod,
					proof.f[j])
			} else {
				fprod = group.Scalar().Mul(fprod,
					group.Scalar().Sub(x,
						proof.f[j]))
			}
			j++
			mask *= 2
		}
		if i < N {
			lhs = group.Point().Add(lhs,
				group.Point().Mul(fprod, c[i]))
		} else {
			lhs = group.Point().Add(lhs,
				group.Point().Mul(fprod, params.Y))
		}
	}
	rhs := group.Point().Mul(proof.zd, params.B)
	k := uint32(0)
	xk := group.Scalar().One() // xk = x^k
	for ; k < n ; {
		rhs = group.Point().Add(rhs,
			group.Point().Mul(xk, pub.cd[k]))
		k++
		xk = group.Scalar().Mul(xk, x)
	}
	if !lhs.Equal(rhs) {
		fmt.Printf("Failed equality check\n")
		return false
	}

	return true
}