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

package stats

import (
	"math"
	"testing"
)

// TestGammaIncompleteClosed checks the incomplete gamma against
// closed forms: P(1, x) = 1 − e^{−x}, P(0.5, x) = erf(√x),
// Q(2, 1) = 3/e and the complement identity.
func TestGammaIncompleteClosed(t *testing.T) {
	p, err := GammaLower(1, 2.5)
	if err != nil {
		t.Fatalf("GammaLower: %v", err)
	}
	if math.Abs(p-(1-math.Exp(-2.5))) > 1e-14 {
		t.Fatalf("P(1, 2.5) = %.16g, want %.16g", p, 1-math.Exp(-2.5))
	}
	p, err = GammaLower(0.5, 1)
	if err != nil {
		t.Fatalf("GammaLower: %v", err)
	}
	if math.Abs(p-math.Erf(1)) > 1e-14 {
		t.Fatalf("P(0.5, 1) = %.16g, want erf(1) = %.16g", p, math.Erf(1))
	}
	q, err := GammaUpper(2, 1)
	if err != nil {
		t.Fatalf("GammaUpper: %v", err)
	}
	if math.Abs(q-2/math.E) > 1e-14 {
		t.Fatalf("Q(2, 1) = %.16g, want 2/e = %.16g", q, 2/math.E)
	}
	p, err = GammaLower(3, 0.7)
	if err != nil {
		t.Fatalf("GammaLower: %v", err)
	}
	q, err = GammaUpper(3, 0.7)
	if err != nil {
		t.Fatalf("GammaUpper: %v", err)
	}
	if math.Abs(p+q-1) > 1e-14 {
		t.Fatalf("P + Q = %.17g, want 1", p+q)
	}
	// The series side of a large shape and the fraction side of a
	// small one must agree with the complement to rounding level.
	p, _ = GammaLower(8.5, 9.3)
	q, _ = GammaUpper(8.5, 9.3)
	if math.Abs(p+q-1) > 1e-13 {
		t.Fatalf("P(8.5, 9.3) + Q(8.5, 9.3) = %.17g, want 1", p+q)
	}
	if _, err := GammaLower(0, 1); err == nil {
		t.Fatal("a = 0: want an error")
	}
	if _, err := GammaLower(1, -1); err == nil {
		t.Fatal("negative x: want an error")
	}
}

// TestBetaIncompleteClosed checks the incomplete beta against closed
// forms and the symmetry I_x(a,b) = 1 − I_{1−x}(b,a).
func TestBetaIncompleteClosed(t *testing.T) {
	cases := []struct {
		x, a, b, want float64
	}{
		{0.3, 1, 1, 0.3},     // uniform
		{0.2, 1, 3, 0.488},   // 1 − (1−x)^b
		{0.09, 2, 1, 0.0081}, // x^a
		{0.5, 2, 2, 0.5},     // symmetric
		{0.5, 0.5, 0.5, 0.5}, // arcsine, symmetric
		{0.25, 0.5, 1, 0.5},  // I_x(1/2,1) = √x
	}
	for _, c := range cases {
		got, err := BetaIncomplete(c.x, c.a, c.b)
		if err != nil {
			t.Fatalf("BetaIncomplete(%g, %g, %g): %v", c.x, c.a, c.b, err)
		}
		if math.Abs(got-c.want) > 1e-13 {
			t.Fatalf("I_%g(%g, %g) = %.16g, want %.16g", c.x, c.a, c.b, got, c.want)
		}
	}
	got, err := BetaIncomplete(0.7, 2, 3)
	if err != nil {
		t.Fatalf("BetaIncomplete: %v", err)
	}
	mirror, err := BetaIncomplete(0.3, 3, 2)
	if err != nil {
		t.Fatalf("BetaIncomplete mirror: %v", err)
	}
	if math.Abs(got+mirror-1) > 1e-13 {
		t.Fatalf("symmetry broken: %.17g + %.17g != 1", got, mirror)
	}
	if _, err := BetaIncomplete(1.5, 1, 1); err == nil {
		t.Fatal("x outside [0, 1]: want an error")
	}
}

// TestContinuousCDFs pins each continuous CDF on exact or tabulated
// values.
func TestContinuousCDFs(t *testing.T) {
	if v := NormalCDF(1); math.Abs(v-0.8413447460685429) > 1e-15 {
		t.Fatalf("Φ(1) = %.16g", v)
	}
	if v := NormalCDF(1.959963984540054); math.Abs(v-0.975) > 1e-14 {
		t.Fatalf("Φ(1.96…) = %.16g, want 0.975", v)
	}
	v, err := ExponentialCDF(1, 1)
	if err != nil || math.Abs(v-(1-1/math.E)) > 1e-15 {
		t.Fatalf("ExponentialCDF(1, 1) = %v, %v", v, err)
	}
	// Gamma(2, 1) CDF: 1 − e^{−x}(1 + x).
	v, err = GammaCDF(1, 2, 1)
	if err != nil || math.Abs(v-(1-2/math.E)) > 1e-14 {
		t.Fatalf("GammaCDF(1, 2, 1) = %v, %v", v, err)
	}
	// χ² with df 2 is the exponential with mean 2: 1 − e^{−x/2}.
	v, err = ChiSquareCDF(1, 2)
	if err != nil || math.Abs(v-(1-math.Exp(-0.5))) > 1e-14 {
		t.Fatalf("ChiSquareCDF(1, 2) = %v, %v", v, err)
	}
	// χ² 95 % critical value with df 1: 3.841458820694124.
	v, err = ChiSquareCDF(3.841458820694124, 1)
	if err != nil || math.Abs(v-0.95) > 1e-13 {
		t.Fatalf("ChiSquareCDF at the tabulated 95 %% point = %v, %v", v, err)
	}
	// Student t with df 1 is the Cauchy: 0.5 + atan(t)/π.
	v, err = StudentTCDF(1, 1)
	if err != nil || math.Abs(v-(0.5+math.Atan(1)/math.Pi)) > 1e-14 {
		t.Fatalf("StudentTCDF(1, 1) = %v, %v", v, err)
	}
	// Student t with df 2: 0.5 + t/(2√(2 + t²)).
	v, err = StudentTCDF(1, 2)
	if err != nil || math.Abs(v-(0.5+1/(2*math.Sqrt(3)))) > 1e-14 {
		t.Fatalf("StudentTCDF(1, 2) = %v, %v", v, err)
	}
}

// TestDiscreteCDFs pins the Poisson and binomial CDFs on exact sums.
func TestDiscreteCDFs(t *testing.T) {
	v, err := PoissonCDF(1, 1)
	if err != nil || math.Abs(v-2/math.E) > 1e-14 {
		t.Fatalf("PoissonCDF(1, 1) = %v, %v, want 2/e", v, err)
	}
	// P(N ≤ 4) for λ = 4: e^{−4}·Σ_{k≤4} 4^k/k!.
	sum := 0.0
	for k, term := 0, 1.0; k <= 4; k++ {
		if k > 0 {
			term *= 4 / float64(k)
		}
		sum += term
	}
	v, err = PoissonCDF(4, 4)
	if err != nil || math.Abs(v-sum*math.Exp(-4)) > 1e-14 {
		t.Fatalf("PoissonCDF(4, 4) = %v, %v, want %.16g", v, err, sum*math.Exp(-4))
	}
	// Binomial(10, 0.5) at 5: 638/1024 by symmetry of the row.
	v, err = BinomialCDF(5, 10, 0.5)
	if err != nil || math.Abs(v-638.0/1024) > 1e-14 {
		t.Fatalf("BinomialCDF(5, 10, 0.5) = %v, %v, want 0.623046875", v, err)
	}
	v, err = BinomialCDF(-1, 10, 0.5)
	if err != nil || v != 0 {
		t.Fatalf("BinomialCDF(-1, …) = %v, %v", v, err)
	}
	v, err = BinomialCDF(10, 10, 0.5)
	if err != nil || v != 1 {
		t.Fatalf("BinomialCDF(10, 10, 0.5) = %v, %v", v, err)
	}
}

// TestContinuousQuantiles inverts every continuous CDF and pins the
// tabulated critical values.
func TestContinuousQuantiles(t *testing.T) {
	v, err := NormalQuantile(0.975)
	if err != nil || math.Abs(v-1.959963984540054) > 1e-12 {
		t.Fatalf("NormalQuantile(0.975) = %v, %v", v, err)
	}
	v, err = StudentTQuantile(0.975, 10)
	if err != nil || math.Abs(v-2.228138851986273) > 1e-9 {
		t.Fatalf("StudentTQuantile(0.975, 10) = %v, %v, want 2.2281…", v, err)
	}
	neg, err := StudentTQuantile(0.025, 10)
	if err != nil || math.Abs(neg+2.228138851986273) > 1e-9 {
		t.Fatalf("StudentTQuantile(0.025, 10) = %v, %v", neg, err)
	}
	v, err = ChiSquareQuantile(0.95, 1)
	if err != nil || math.Abs(v-3.841458820694124) > 1e-9 {
		t.Fatalf("ChiSquareQuantile(0.95, 1) = %v, %v", v, err)
	}
	v, err = ExponentialQuantile(0.6321205588285577, 1)
	if err != nil || math.Abs(v-1) > 1e-12 {
		t.Fatalf("ExponentialQuantile(1−1/e, 1) = %v, %v", v, err)
	}
	v, err = GammaQuantile(0.5, 2, 1)
	if err != nil || math.Abs(v-1.678346990016661) > 1e-9 {
		t.Fatalf("GammaQuantile(0.5, 2, 1) = %v, %v, want 1.67834…", v, err)
	}
	// Round trip: the CDF at every quantile must return q.
	for _, q := range []float64{0.01, 0.1, 0.5, 0.9, 0.999} {
		v, err := GammaQuantile(q, 3.5, 2)
		if err != nil {
			t.Fatalf("GammaQuantile(%g): %v", q, err)
		}
		back, err := GammaCDF(v, 3.5, 2)
		if err != nil || math.Abs(back-q) > 1e-11 {
			t.Fatalf("round trip q = %g: CDF(quantile) = %v, %v", q, back, err)
		}
	}
	if _, err := NormalQuantile(1.5); err == nil {
		t.Fatal("q outside [0, 1]: want an error")
	}
	if _, err := NormalQuantile(0); err == nil {
		t.Fatal("q = 0: want an error")
	}
}

// TestDiscreteQuantiles checks the smallest-k rule on exact cases.
func TestDiscreteQuantiles(t *testing.T) {
	v, err := PoissonQuantile(0.5, 1)
	if err != nil || v != 1 {
		t.Fatalf("PoissonQuantile(0.5, 1) = %v, %v, want 1", v, err)
	}
	// P(N ≤ 2) = 5/(2e) ≈ 0.9197, P(N ≤ 1) = 2/e ≈ 0.7358 for λ = 1:
	// the 0.8 quantile is the smallest k reaching it, k = 2.
	v, err = PoissonQuantile(0.8, 1)
	if err != nil || v != 2 {
		t.Fatalf("PoissonQuantile(0.8, 1) = %v, %v, want 2", v, err)
	}
	// Binomial(10, 0.5) median: smallest k with CDF ≥ 0.5 is 5.
	v, err = BinomialQuantile(0.5, 0.5, 10)
	if err != nil || v != 5 {
		t.Fatalf("BinomialQuantile(0.5, 0.5, 10) = %v, %v, want 5", v, err)
	}
	if _, err := PoissonQuantile(0.5, 0); err == nil {
		t.Fatal("lambda = 0: want an error")
	}
}

// TestGammaLowerLargeShape pins the large-shape region: both the power
// series (x < a+1) and the continued fraction (x ≥ a+1) must deliver
// accurate values near x ≈ a where the shape makes √a-sized iteration
// counts necessary, instead of silently returning truncated sums.
// The reference is the Wilson-Hilferty normal approximation, good to
// a few digits at these shapes.
func TestGammaLowerLargeShape(t *testing.T) {
	wh := func(a, x float64) float64 {
		z := 3 * math.Sqrt(a) * (math.Cbrt(x/a) - (1 - 1/(9*a)))
		return 0.5 * math.Erfc(-z/math.Sqrt2)
	}
	cases := []struct {
		a, x, tol float64
	}{
		{50000, 50000, 1e-4}, // series branch, √a ≈ 224
		{50000, 49500, 5e-4}, // lower tail, series
		{50000, 50500, 5e-4}, // upper tail, continued fraction
		{200000, 200000, 1e-4},
		{1000, 1000, 1e-5},
	}
	for _, tc := range cases {
		p, err := GammaLower(tc.a, tc.x)
		if err != nil {
			t.Fatalf("GammaLower(%g, %g): %v", tc.a, tc.x, err)
		}
		want := wh(tc.a, tc.x)
		if math.Abs(p-want) > tc.tol {
			t.Fatalf("GammaLower(%g, %g) = %.10g, want ≈ %.10g (tol %g)", tc.a, tc.x, p, want, tc.tol)
		}
		q, err := GammaUpper(tc.a, tc.x)
		if err != nil {
			t.Fatalf("GammaUpper(%g, %g): %v", tc.a, tc.x, err)
		}
		if p+q != 1 {
			t.Fatalf("GammaLower + GammaUpper = %.17g, want exactly 1", p+q)
		}
	}
}

// TestStudentTCDFSquareOverflowTail pins the Student laws past the square
// overflow: a t beyond √MaxFloat64 drives z = df/(df+t²) through Inf/Inf
// to the NaN and 0 the closed form answers, while the heavy df = 1 and
// df = 2 tails are still representable there. The lower tail and the
// two-sided tail must answer their asymptotic forms, exactly as the
// one-sided upper tail already does.
func TestStudentTCDFSquareOverflowTail(t *testing.T) {
	const huge = 1e155
	// df = 1 (Cauchy): the one-sided tail is 1/(π·t), two-sided 2/(π·t).
	one, err := StudentTCDF(-huge, 1)
	if err != nil {
		t.Fatalf("StudentTCDF(-1e155, 1): %v", err)
	}
	if want := 1 / (math.Pi * huge); math.Abs(one-want) > 1e-12*want {
		t.Fatalf("StudentTCDF(-1e155, 1) = %.17g, want the Cauchy tail %.17g", one, want)
	}
	two, err := twoSidedT(huge, 1)
	if err != nil {
		t.Fatalf("twoSidedT(1e155, 1): %v", err)
	}
	if want := 2 / (math.Pi * huge); math.Abs(two-want) > 1e-12*want {
		t.Fatalf("twoSidedT(1e155, 1) = %.17g, want the Cauchy tail %.17g", two, want)
	}
	// df = 2: the two-sided tail is 1 − t/√(t²+2), the one-sided half of
	// it, both ≈ 1/t² here and still inside the subnormal range. The
	// reference is assembled as u/(√(1+u)+1) with u = 2/t², the form
	// that never squares t.
	const u = 2e-310 // 2/t² at t = 1e155
	two2, err := twoSidedT(huge, 2)
	if err != nil {
		t.Fatalf("twoSidedT(1e155, 2): %v", err)
	}
	if want := u / (math.Sqrt(1+u) + 1); math.Abs(two2-want) > 1e-6*want {
		t.Fatalf("twoSidedT(1e155, 2) = %.17g, want the df 2 tail %.17g", two2, want)
	}
	one2, err := StudentTCDF(-huge, 2)
	if err != nil {
		t.Fatalf("StudentTCDF(-1e155, 2): %v", err)
	}
	if want := 0.5 * u / (math.Sqrt(1+u) + 1); math.Abs(one2-want) > 1e-6*want {
		t.Fatalf("StudentTCDF(-1e155, 2) = %.17g, want the df 2 tail %.17g", one2, want)
	}
	// The upper half of the axis keeps answering 1, and a df whose tail
	// genuinely underflows keeps answering 0: both are the honest
	// roundings there.
	if v, err := StudentTCDF(huge, 1); err != nil || v != 1 {
		t.Fatalf("StudentTCDF(1e155, 1) = %v (%v), want 1", v, err)
	}
	if v, err := StudentTCDF(-huge, 4); err != nil || v != 0 {
		t.Fatalf("StudentTCDF(-1e155, 4) = %v (%v), want the underflowed 0", v, err)
	}
}