// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package stats

import (
	"math"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
	"testing"
)

// TestDistributions checks the seeded draws against theoretical
// moments with generous statistical tolerances, and the error
// contract for invalid parameters.
func TestDistributions(t *testing.T) {
	g := core.NewGenerator(42)
	const n = 200000

	// Exponential(rate 2): mean ½, variance ¼.
	exp, err := ExponentialDraws(g, n, 2)
	if err != nil {
		t.Fatalf("ExponentialDraws: %v", err)
	}
	mean := core.Sum(exp).Float()
	if math.Abs(mean/float64(n)-0.5) > 0.01 {
		t.Fatalf("exponential mean = %v, want ≈ 0.5", mean/float64(n))
	}
	for i := range n {
		if exp.FloatAt(i) < 0 {
			t.Fatalf("exponential draw %d = %v, must be positive", i, exp.FloatAt(i))
		}
	}

	// Gamma(3, 2): mean 3/2, variance 3/4.
	gam, err := GammaDraws(g, n, 3, 2)
	if err != nil {
		t.Fatalf("GammaDraws: %v", err)
	}
	gsum := core.Sum(gam).Float()
	if math.Abs(gsum/float64(n)-1.5) > 0.03 {
		t.Fatalf("gamma mean = %v, want ≈ 1.5", gsum/float64(n))
	}

	// ChiSquare(5): mean 5.
	chi, _ := ChiSquareDraws(g, n, 5)
	csum := core.Sum(chi).Float()
	if math.Abs(csum/float64(n)-5) > 0.1 {
		t.Fatalf("chi² mean = %v, want ≈ 5", csum/float64(n))
	}

	// Poisson(4): mean 4, variance 4.
	pois, _ := PoissonDraws(g, n, 4)
	psum := core.Sum(pois).Float()
	if math.Abs(psum/float64(n)-4) > 0.1 {
		t.Fatalf("poisson mean = %v, want ≈ 4", psum/float64(n))
	}

	// Binomial(20, 0.3): mean 6.
	bin, _ := BinomialDraws(g, n, 20, 0.3)
	bsum := core.Sum(bin).Float()
	if math.Abs(bsum/float64(n)-6) > 0.1 {
		t.Fatalf("binomial mean = %v, want ≈ 6", bsum/float64(n))
	}

	// StudentT(5): mean 0.
	st, _ := StudentTDraws(g, n, 5)
	ssum := core.Sum(st).Float()
	if math.Abs(ssum/float64(n)) > 0.05 {
		t.Fatalf("t mean = %v, want ≈ 0", ssum/float64(n))
	}
}

// TestDistributionsErrors pins the parameter contracts.
func TestDistributionsErrors(t *testing.T) {
	g := core.NewGenerator(1)
	if _, err := ExponentialDraws(g, 1, -1); err == nil {
		t.Fatal("expected an error for a negative rate")
	}
	if _, err := GammaDraws(g, 1, 0, 1); err == nil {
		t.Fatal("expected an error for a zero shape")
	}
	if _, err := ChiSquareDraws(g, 1, 0); err == nil {
		t.Fatal("expected an error for df = 0")
	}
	if _, err := StudentTDraws(g, 1, -1); err == nil {
		t.Fatal("expected an error for negative df")
	}
	if _, err := PoissonDraws(g, 1, -1); err == nil {
		t.Fatal("expected an error for a negative λ")
	}
	if _, err := BinomialDraws(g, 1, 10, 1.5); err == nil {
		t.Fatal("expected an error for p > 1")
	}
}