package gnarlring

import (
	"math"
	"testing"
)

func TestCastVote(t *testing.T) {
	yes := CastVote(1)
	no := CastVote(-1)

	if IsZero(yes.Z) || IsZero(no.Z) {
		t.Fatal("vote vectors are zero")
	}

	// YES: first coefficient = 1. NO: first coefficient = Q-1.
	if yes.Z.Coeffs[0] != 1 {
		t.Fatalf("YES vote coeff[0] = %d, want 1", yes.Z.Coeffs[0])
	}
	if no.Z.Coeffs[0] != Q-1 {
		t.Fatalf("NO vote coeff[0] = %d, want %d", no.Z.Coeffs[0], Q-1)
	}

	// Other coefficients should be Gaussian (non-zero on average).
	nonZeroCount := 0
	for i := 1; i < N; i++ {
		if yes.Z.Coeffs[i] != 0 || no.Z.Coeffs[i] != 0 {
			nonZeroCount++
		}
	}
	if nonZeroCount < 10 {
		t.Fatalf("only %d/26 non-zero Gaussian coefficients — unlikely", nonZeroCount)
	}
}

func TestVoteTally(t *testing.T) {
	tally := NewVoteTally()
	for i := 0; i < 15; i++ {
		tally.Add(CastVote(1), true)
	}
	for i := 0; i < 12; i++ {
		tally.Add(CastVote(-1), false)
	}

	if tally.YesCount != 15 || tally.NoCount != 12 {
		t.Fatalf("yes=%d no=%d", tally.YesCount, tally.NoCount)
	}

	if !tally.ConsensusResult(10) {
		t.Fatal("threshold 10 should pass with 15 yes")
	}
	if !tally.ConsensusResult(14) {
		t.Fatal("threshold 14 should pass with 15 yes")
	}
	if tally.ConsensusResult(16) {
		t.Fatal("threshold 16 should fail with 15 yes")
	}

	normVal := math.Sqrt(float64(NormSq(tally.Z)))
	maxPossible := float64(tally.YesCount+tally.NoCount) * 150.0
	t.Logf("tally norm=%.0f max=%.0f", normVal, maxPossible)
	if normVal > maxPossible {
		t.Fatalf("tally norm %.0f exceeds max %.0f", normVal, maxPossible)
	}
}

func TestEncryptedVote(t *testing.T) {
	lwePK, lweSK := LWEKeyGen()

	yes := CastVote(1)
	ev := EncryptVote(lwePK, yes, nil)
	if ev.Ct == nil {
		t.Fatal("encrypted vote is nil")
	}

	result := DecryptVote(lweSK, ev)
	if !result {
		t.Fatal("decrypted vote should be YES")
	}

	no := CastVote(-1)
	evNo := EncryptVote(lwePK, no, nil)
	resultNo := DecryptVote(lweSK, evNo)
	if resultNo {
		t.Fatal("decrypted NO vote should be false")
	}

	correct := 0
	for trial := 0; trial < 50; trial++ {
		v := CastVote(1)
		ev := EncryptVote(lwePK, v, nil)
		if DecryptVote(lweSK, ev) {
			correct++
		}
	}
	t.Logf("encrypted vote accuracy: %d/50", correct)
	if correct < 45 {
		t.Fatalf("encrypted vote accuracy %d/50 < 90%%", correct)
	}
}

func BenchmarkCastVote(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		CastVote(1)
	}
}

func BenchmarkVoteSum(b *testing.B) {
	votes := make([]*Vote, 27)
	for i := range votes {
		votes[i] = CastVote(1)
	}
	tally := NewVoteTally()
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		for _, v := range votes {
			tally.Add(v, true)
		}
	}
}