package crypto

import (
	"math/rand"
	"testing"
)

func TestDigitRev3(t *testing.T) {
	// For 3 digits (n=27), digit reversal should be an involution.
	for i := range 27 {
		j := digitRev3(i, 3)
		if j < 0 || j >= 27 {
			t.Fatalf("digitRev3(%d, 3) = %d out of range", i, j)
		}
		k := digitRev3(j, 3)
		if k != i {
			t.Fatalf("digitRev3 not involution: digitRev3(%d) = %d, digitRev3(%d) = %d", i, j, j, k)
		}
	}

	// Check specific values.
	// 0 in base 3 = 000, reversed = 000 = 0
	if digitRev3(0, 3) != 0 {
		t.Fatalf("digitRev3(0, 3) = %d, want 0", digitRev3(0, 3))
	}
	// 1 in base 3 = 001, reversed = 100 = 9
	if digitRev3(1, 3) != 9 {
		t.Fatalf("digitRev3(1, 3) = %d, want 9", digitRev3(1, 3))
	}
	// 3 in base 3 = 010, reversed = 010 = 3
	if digitRev3(3, 3) != 3 {
		t.Fatalf("digitRev3(3, 3) = %d, want 3", digitRev3(3, 3))
	}
}

func TestNTT27RoundTrip(t *testing.T) {
	rng := rand.New(rand.NewSource(42))

	for trial := range 10 {
		var a [GnarlN]uint16
		for i := range GnarlN {
			a[i] = uint16(rng.Intn(GnarlP))
		}

		orig := a
		ntt27(&a)
		intt27(&a)

		for i := range GnarlN {
			if a[i] != orig[i] {
				t.Fatalf("trial %d: round-trip failed at index %d: got %d, want %d",
					trial, i, a[i], orig[i])
			}
		}
	}
}

func TestNTT27Zero(t *testing.T) {
	var a [GnarlN]uint16
	ntt27(&a)
	for i, v := range a {
		if v != 0 {
			t.Fatalf("NTT of zero polynomial: index %d = %d, want 0", i, v)
		}
	}
	intt27(&a)
	for i, v := range a {
		if v != 0 {
			t.Fatalf("INTT of zero: index %d = %d, want 0", i, v)
		}
	}
}

func TestNTT27Convolution(t *testing.T) {
	// Pointwise multiply in NTT domain should equal polynomial multiplication
	// modulo (x^27 + 1) in coefficient domain.
	rng := rand.New(rand.NewSource(99))

	var a, b [GnarlN]uint16
	for i := range GnarlN {
		a[i] = uint16(rng.Intn(GnarlP))
		b[i] = uint16(rng.Intn(GnarlP))
	}

	// Schoolbook multiplication mod (x^27 + 1).
	var expected [GnarlN]uint16
	for i := range GnarlN {
		for j := range GnarlN {
			k := i + j
			if k < GnarlN {
				expected[k] = mod271(uint32(expected[k]) + uint32(a[i])*uint32(b[j]))
			} else {
				// x^27 = -1 mod (x^27 + 1), so x^(27+r) = -x^r
				k -= GnarlN
				expected[k] = mod271s(int(expected[k]) - int(a[i])*int(b[j]))
			}
		}
	}

	// NTT multiplication.
	aNtt := a
	bNtt := b
	ntt27(&aNtt)
	ntt27(&bNtt)

	var cNtt [GnarlN]uint16
	for i := range GnarlN {
		cNtt[i] = mod271(uint32(aNtt[i]) * uint32(bNtt[i]))
	}
	intt27(&cNtt)

	for i := range GnarlN {
		if cNtt[i] != expected[i] {
			t.Fatalf("convolution mismatch at index %d: NTT gives %d, schoolbook gives %d",
				i, cNtt[i], expected[i])
		}
	}
}

func TestNTT27Negacyclic(t *testing.T) {
	// x^27 ≡ -1 mod (x^27 + 1).
	// Represent x^27 as polynomial: coefficient 0 at index 27, which wraps to
	// -1 * x^0. So the polynomial [p-1, 0, 0, ...] should be the product of
	// [0, 1, 0, ...] raised to the 27th power.
	//
	// Simpler test: multiply [0,1,0,...,0] (which is x) by itself 27 times
	// and check we get [-1, 0, ..., 0].
	//
	// Actually, just multiply x^26 * x = x^27. Start with x, multiply by x
	// in NTT domain 26 times.

	// x = [0, 1, 0, ..., 0]
	var x [GnarlN]uint16
	x[1] = 1

	// x in NTT domain.
	xNtt := x
	ntt27(&xNtt)

	// x^27 in NTT domain = pointwise x^27.
	var result [GnarlN]uint16
	for i := range GnarlN {
		result[i] = powMod271(int(xNtt[i]), 27)
	}
	intt27(&result)

	// Should be [-1, 0, 0, ..., 0] = [270, 0, 0, ..., 0].
	if result[0] != uint16(GnarlP-1) {
		t.Fatalf("x^27 coefficient 0: got %d, want %d", result[0], GnarlP-1)
	}
	for i := 1; i < GnarlN; i++ {
		if result[i] != 0 {
			t.Fatalf("x^27 coefficient %d: got %d, want 0", i, result[i])
		}
	}
}

func TestNTT27PsiPowers(t *testing.T) {
	// Verify psi^54 = 1 and psi^27 = p-1 = -1.
	if psiPows27[0] != 1 {
		t.Fatalf("psi^0 = %d, want 1", psiPows27[0])
	}
	// psi^27 should be p-1 (since psi is a 54th root, psi^27 = psi^(54/2) = -1).
	if psiPows27[27] != uint16(GnarlP-1) {
		t.Fatalf("psi^27 = %d, want %d", psiPows27[27], GnarlP-1)
	}

	// Verify invN27.
	prod := mod271(uint32(invN27) * GnarlN)
	if prod != 1 {
		t.Fatalf("invN27 * 27 mod 271 = %d, want 1", prod)
	}
}

func BenchmarkNTT27(b *testing.B) {
	var a [GnarlN]uint16
	for i := range GnarlN {
		a[i] = uint16(i * 10 % GnarlP)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ntt27(&a)
	}
}

func BenchmarkINTT27(b *testing.B) {
	var a [GnarlN]uint16
	for i := range GnarlN {
		a[i] = uint16(i * 10 % GnarlP)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		intt27(&a)
	}
}

func BenchmarkMod271(b *testing.B) {
	x := uint32(50000)
	for i := 0; i < b.N; i++ {
		_ = mod271(x)
		x = uint32(mod271(x))*270 + 1
	}
}

func BenchmarkGnarlCompress(b *testing.B) {
	var block [gnarlInputBytes]byte
	for i := range block {
		block[i] = byte(i * 37)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		gnarlCompress(&block)
	}
}

func TestNTT27Linearity(t *testing.T) {
	// NTT should be linear: NTT(a + b) = NTT(a) + NTT(b).
	rng := rand.New(rand.NewSource(777))

	var a, b, sum [GnarlN]uint16
	for i := range GnarlN {
		a[i] = uint16(rng.Intn(GnarlP))
		b[i] = uint16(rng.Intn(GnarlP))
		sum[i] = mod271(uint32(a[i]) + uint32(b[i]))
	}

	ntt27(&a)
	ntt27(&b)
	ntt27(&sum)

	for i := range GnarlN {
		expected := mod271(uint32(a[i]) + uint32(b[i]))
		if sum[i] != expected {
			t.Fatalf("linearity failed at index %d: NTT(a+b)[%d]=%d, NTT(a)+NTT(b)=%d",
				i, i, sum[i], expected)
		}
	}
}