package ring

import (
	"math"
	"testing"
)

// TestGaussianSampleZDistribution verifies that SampleZ produces
// a distribution that looks Gaussian: concentrated near 0 with the
// right spread.
func TestGaussianSampleZDistribution(t *testing.T) {
	sigma := 5.0
	gs := NewGaussianSampler(sigma)

	n := 100000
	samples := make([]int64, n)
	for i := range samples {
		samples[i] = gs.SampleZ(0)
	}

	// Compute mean and variance.
	sum := 0.0
	for _, s := range samples {
		sum += float64(s)
	}
	mean := sum / float64(n)

	sumSq := 0.0
	for _, s := range samples {
		d := float64(s) - mean
		sumSq += d * d
	}
	variance := sumSq / float64(n)

	// Expected variance for D_{Z, σ, 0} ≈ σ²/(2π).
	expectedVar := sigma * sigma / (2 * math.Pi)

	t.Logf("σ=%.1f: mean=%.3f, variance=%.3f (expected ≈ %.3f)",
		sigma, mean, variance, expectedVar)

	// Mean should be near 0.
	if math.Abs(mean) > 0.5 {
		t.Errorf("mean too far from 0: %.3f", mean)
	}

	// Variance should be within 20% of expected.
	ratio := variance / expectedVar
	if ratio < 0.7 || ratio > 1.3 {
		t.Errorf("variance ratio %.3f outside [0.7, 1.3]", ratio)
	}
}

// TestGaussianSampleZCentered verifies centered sampling.
func TestGaussianSampleZCentered(t *testing.T) {
	sigma := 8.0
	center := 42.0
	gs := NewGaussianSampler(sigma)

	n := 50000
	sum := 0.0
	for range n {
		s := gs.SampleZ(center)
		sum += float64(s)
	}
	mean := sum / float64(n)

	t.Logf("σ=%.1f, c=%.1f: mean=%.3f", sigma, center, mean)

	if math.Abs(mean-center) > 1.0 {
		t.Errorf("mean %.3f too far from center %.1f", mean, center)
	}
}

// TestGaussianSampleZLargeSigma verifies the rejection sampler
// for larger σ values used in GPV.
func TestGaussianSampleZLargeSigma(t *testing.T) {
	sigma := 50.0
	gs := NewGaussianSampler(sigma)

	n := 50000
	sum := 0.0
	maxAbs := int64(0)
	for range n {
		s := gs.SampleZ(0)
		sum += float64(s)
		if s < 0 {
			s = -s
		}
		if s > maxAbs {
			maxAbs = s
		}
	}
	mean := sum / float64(n)

	t.Logf("σ=%.1f: mean=%.3f, max|sample|=%d (tail bound=%.0f)",
		sigma, mean, maxAbs, 13.0*sigma)

	// Max should be within tail bound.
	if float64(maxAbs) > 13.0*sigma {
		t.Errorf("sample %d exceeds tail bound %.0f", maxAbs, 13.0*sigma)
	}

	if math.Abs(mean) > 2.0 {
		t.Errorf("mean %.3f too far from 0", mean)
	}
}

// TestGaussianSamplePoly verifies polynomial sampling.
func TestGaussianSamplePoly(t *testing.T) {
	sigma := 5.0
	gs := NewGaussianSampler(sigma)
	p := Falcon512()

	poly := gs.SamplePoly(p)

	// All coefficients should be in [0, Q).
	for i, c := range poly.Coeffs {
		if c >= p.Q {
			t.Fatalf("coeff[%d] = %d >= Q=%d", i, c, p.Q)
		}
	}

	// Count small coefficients (centered representation).
	// Most should be near 0 (i.e., small or near Q).
	smallCount := 0
	half := p.Q / 2
	for _, c := range poly.Coeffs {
		var v uint32
		if c > half {
			v = p.Q - c
		} else {
			v = c
		}
		if v < 20 {
			smallCount++
		}
	}

	fraction := float64(smallCount) / float64(p.N)
	t.Logf("σ=%.1f: %.1f%% of coefficients within ±20", sigma, fraction*100)

	// With σ=5, most coefficients should be small.
	if fraction < 0.5 {
		t.Errorf("too few small coefficients: %.1f%%", fraction*100)
	}
}

// TestGaussianCDTvRejection verifies CDT and rejection produce
// similar distributions.
func TestGaussianCDTvRejection(t *testing.T) {
	sigma := 5.0

	gsCDT := NewGaussianSampler(sigma)
	if gsCDT.cdt == nil {
		t.Fatal("CDT not built for σ=5")
	}

	// Also create a rejection sampler by using larger sigma threshold.
	gsRej := &GaussianSampler{
		Sigma:     sigma,
		TailBound: 13.0,
		rng:       gsCDT.rng,
	}
	// Force rejection path by not building CDT.

	n := 50000

	// Build histograms.
	histCDT := make(map[int64]int)
	histRej := make(map[int64]int)

	for range n {
		s := gsCDT.sampleCDT(0)
		histCDT[s]++
	}
	for range n {
		s := gsRej.sampleRejection(0)
		histRej[s]++
	}

	// Compare distributions using chi-squared-like test.
	// Just verify both are peaked at 0 and symmetric.
	if histCDT[0] < n/10 {
		t.Errorf("CDT: too few zeros: %d/%d", histCDT[0], n)
	}
	if histRej[0] < n/10 {
		t.Errorf("Rejection: too few zeros: %d/%d", histRej[0], n)
	}

	t.Logf("CDT[0]=%d, CDT[±1]=%d/%d, CDT[±2]=%d/%d",
		histCDT[0], histCDT[1], histCDT[-1], histCDT[2], histCDT[-2])
	t.Logf("Rej[0]=%d, Rej[±1]=%d/%d, Rej[±2]=%d/%d",
		histRej[0], histRej[1], histRej[-1], histRej[2], histRej[-2])
}

// TestRingGaussianSigma verifies the recommended sigma values.
func TestRingGaussianSigma(t *testing.T) {
	for _, n := range []int{64, 256, 512, 1024} {
		sigma := RingGaussianSigma(n)
		t.Logf("n=%d: σ=%.1f", n, sigma)
		if sigma <= 0 {
			t.Errorf("n=%d: σ=%.1f should be positive", n, sigma)
		}
	}
}

func BenchmarkGaussianSampleZ(b *testing.B) {
	gs := NewGaussianSampler(5.0)
	b.ResetTimer()
	for range b.N {
		gs.SampleZ(0)
	}
}

func BenchmarkGaussianSampleZLarge(b *testing.B) {
	gs := NewGaussianSampler(165.0)
	b.ResetTimer()
	for range b.N {
		gs.SampleZ(0)
	}
}

func BenchmarkGaussianSamplePoly(b *testing.B) {
	gs := NewGaussianSampler(165.0)
	p := Falcon512()
	b.ResetTimer()
	for range b.N {
		gs.SamplePoly(p)
	}
}