// 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"
	"strings"
	"testing"
)

// TestNoncentralChiSquareClosed pins the noncentral χ² on the exact
// closed form its df = 1 corner carries: χ²(1, λ) is the square of a
// N(√λ, 1) draw, so the CDF is Φ(√x−√λ) − Φ(−√x−√λ), and on the λ = 0
// reduction to the central law.
func TestNoncentralChiSquareClosed(t *testing.T) {
	for _, lambda := range []float64{0.5, 1, 4, 9, 25} {
		for _, x := range []float64{0.5, 1, 2, 4, 9, 16} {
			got, err := NoncentralChiSquareCDF(x, 1, lambda)
			if err != nil {
				t.Fatalf("NoncentralChiSquareCDF(%g, 1, %g): %v", x, lambda, err)
			}
			root := math.Sqrt(lambda)
			want := NormalCDF(math.Sqrt(x)-root) - NormalCDF(-math.Sqrt(x)-root)
			if math.Abs(got-want) > 1e-13 {
				t.Fatalf("NoncentralChiSquareCDF(%g, 1, %g) = %.16g, want %.16g", x, lambda, got, want)
			}
		}
	}
	for _, df := range []int{1, 2, 5, 10} {
		for _, x := range []float64{0.5, 2, 7} {
			got, err := NoncentralChiSquareCDF(x, df, 0)
			if err != nil {
				t.Fatalf("NoncentralChiSquareCDF(%g, %d, 0): %v", x, df, err)
			}
			want, err := ChiSquareCDF(x, df)
			if err != nil || math.Abs(got-want) > 1e-14 {
				t.Fatalf("λ = 0 reduction at df %d: %v vs %v (%v)", df, got, want, err)
			}
		}
	}
	if v, _ := NoncentralChiSquareCDF(-1, 3, 2); v != 0 {
		t.Fatalf("CDF below the support = %v, want 0", v)
	}
	if v, _ := NoncentralChiSquareDensity(-1, 3, 2); v != 0 {
		t.Fatalf("density below the support = %v, want 0", v)
	}
}

// TestNoncentralChiSquareDensityIntegral integrates the density against
// the CDF. The walk runs after the substitution x = s², which leaves
// 2s·f(s²) smooth at the origin for every df (the density itself
// behaves like x^{df/2−1} there, too flat a start for Simpson's error
// estimate on the lower degrees). Simpson must then reproduce the CDF
// to better than 1e-10 relative, the double route the closed forms
// cannot cover.
func TestNoncentralChiSquareDensityIntegral(t *testing.T) {
	type grid struct {
		df     int
		lambda float64
		x      float64
	}
	for _, g := range []grid{
		{2, 1, 6}, {3, 1, 6}, {3, 4, 10}, {5, 3, 12}, {10, 25, 60},
	} {
		n := 200000
		s := math.Sqrt(g.x)
		sh := s / float64(n)
		f := func(sv float64) float64 {
			if sv == 0 {
				return 0
			}
			xv := sv * sv
			d, err := NoncentralChiSquareDensity(xv, g.df, g.lambda)
			if err != nil {
				t.Fatalf("NoncentralChiSquareDensity: %v", err)
			}
			return 2 * sv * d
		}
		sum := f(0) + f(s)
		for i := 1; i < n; i++ {
			w := 4.0
			if i%2 == 0 {
				w = 2
			}
			sum += w * f(float64(i)*sh)
		}
		integral := sum * sh / 3
		cdf, err := NoncentralChiSquareCDF(g.x, g.df, g.lambda)
		if err != nil {
			t.Fatalf("NoncentralChiSquareCDF: %v", err)
		}
		if rel := math.Abs(integral-cdf) / cdf; rel > 1e-10 {
			t.Fatalf("df = %d, λ = %g, x = %g: density integral %.15g vs CDF %.15g (rel %g)",
				g.df, g.lambda, g.x, integral, cdf, rel)
		}
	}
}

// noncentralTOracle evaluates E[Φ(t√(V/ν) − δ)] for V ~ χ²(ν) by
// Simpson after the substitution V = u², which leaves the integrand
// smooth for every ν; the u = 0 limit is finite only for ν = 1.
func noncentralTOracle(t float64, df int, delta float64) float64 {
	uhi := 40.0
	n := 200000
	h := uhi / float64(n)
	logGammaB := func(a float64) float64 { l, _ := math.Lgamma(a); return l }
	f := func(u float64) float64 {
		if u == 0 {
			if df == 1 {
				return math.Sqrt(2/math.Pi) * NormalCDF(-delta)
			}
			return 0
		}
		v := u * u
		fv := 2 * u * math.Exp((float64(df)/2-1)*math.Log(v)-v/2-(float64(df)/2)*math.Ln2-logGammaB(float64(df)/2))
		return NormalCDF(t*math.Sqrt(v/float64(df))-delta) * fv
	}
	sum := f(0) + f(uhi)
	for i := 1; i < n; i++ {
		w := 4.0
		if i%2 == 0 {
			w = 2
		}
		sum += w * f(float64(i)*h)
	}
	return sum * h / 3
}

// TestNoncentralTAgainstQuadrature holds the Lenth series against a
// direct quadrature of E[Φ(t√(V/ν) − δ)], an independent route that
// shares no code with the series, at a grid spanning both signs of t
// and δ and degrees of freedom from 1 to 10. The large-δ cases are the
// underflow round: a noncentrality whose Poisson weight seed e^{−δ²/2}
// is below the double floor used to silence the whole series, and each
// of them once answered a silent 0. Tolerance 1e-9, an order above the
// quadrature's own accuracy.
func TestNoncentralTAgainstQuadrature(t *testing.T) {
	for _, df := range []int{1, 2, 5, 10} {
		for _, delta := range []float64{-3, -0.5, 0.5, 2} {
			for _, tv := range []float64{-2, -0.5, 0.4, 1, 3} {
				got, err := NoncentralTCDF(tv, df, delta)
				if err != nil {
					t.Fatalf("NoncentralTCDF(%g, %d, %g): %v", tv, df, delta, err)
				}
				want := noncentralTOracle(tv, df, delta)
				if math.Abs(got-want) > 1e-9 {
					t.Fatalf("NoncentralTCDF(%g, %d, %g) = %.15g, want quadrature %.15g",
						tv, df, delta, got, want)
				}
			}
		}
	}
	for _, c := range []struct {
		tv, delta float64
		df        int
	}{
		{45, 45, 5},
		{50, 40, 3},
		{-45, -45, 5},
		{100, 45, 5},
	} {
		got, err := NoncentralTCDF(c.tv, c.df, c.delta)
		if err != nil {
			t.Fatalf("NoncentralTCDF(%g, %d, %g): %v", c.tv, c.df, c.delta, err)
		}
		want := noncentralTOracle(c.tv, c.df, c.delta)
		if math.Abs(got-want) > 1e-9 {
			t.Fatalf("NoncentralTCDF(%g, %d, %g) = %.15g, want quadrature %.15g",
				c.tv, c.df, c.delta, got, want)
		}
	}
}

// TestNoncentralUnderflowSurvival pins the deep-noncentrality window
// where the mixture weights' raw seed underflows: the χ² CDF at its
// own mean answers a half, the far tail answers a genuinely computed
// negligible value rather than a silent zero, the F CDF saturates at
// 1 past the overflow of ν₁x, and the χ² density at the origin keeps
// the finite df = 2 limit. Past the term budget the refusal is
// explicit.
func TestNoncentralUnderflowSurvival(t *testing.T) {
	atMean, err := NoncentralChiSquareCDF(2005, 5, 2000)
	if err != nil {
		t.Fatal(err)
	}
	if atMean < 0.48 || atMean > 0.52 {
		t.Fatalf("NoncentralChiSquareCDF at the mean 2005 = %g, want near a half", atMean)
	}
	tail, err := NoncentralChiSquareCDF(1005, 5, 2000)
	if err != nil {
		t.Fatal(err)
	}
	if !(tail >= 0 && tail < 1e-30) {
		t.Fatalf("NoncentralChiSquareCDF(1005, 5, 2000) = %g, want a negligible non-negative tail", tail)
	}
	for _, x := range []float64{math.MaxFloat64, math.Inf(1)} {
		f, err := NoncentralFCDF(x, 4, 10, 3)
		if err != nil {
			t.Fatal(err)
		}
		if math.Abs(f-1) > 1e-15 {
			t.Fatalf("NoncentralFCDF(%g, 4, 10, 3) = %g, want 1 to rounding", x, f)
		}
	}
	d, err := NoncentralChiSquareDensity(0, 2, 3)
	if err != nil {
		t.Fatal(err)
	}
	if want := 0.5 * math.Exp(-1.5); d != want {
		t.Fatalf("NoncentralChiSquareDensity(0, 2, 3) = %g, want the limit %g", d, want)
	}
	if _, err := NoncentralChiSquareCDF(10, 5, 4e5); err == nil || !strings.Contains(err.Error(), "budget") {
		t.Fatalf("lambda 4e5: error = %v, want the budget refusal", err)
	}
	if _, err := NoncentralTCDF(10, 5, 500); err == nil || !strings.Contains(err.Error(), "budget") {
		t.Fatalf("delta 500: error = %v, want the budget refusal", err)
	}
}

// TestNoncentralTCDFHugeFiniteT pins the far corner of the signed axis: a
// finite t whose square overflows drives the beta argument to Inf/Inf, a
// NaN the incomplete beta refused under its own name. The CDF there is 1
// below rounding for t on the δ side and 0 above it, the same limits the
// central law answers.
func TestNoncentralTCDFHugeFiniteT(t *testing.T) {
	for _, c := range []struct {
		tv    float64
		df    int
		delta float64
		want  float64
	}{
		{1e200, 3, 2, 1},
		{1e155, 1, 0.5, 1},
		{-1e200, 5, 1, 0},
		{-1e155, 2, -3, 0},
	} {
		got, err := NoncentralTCDF(c.tv, c.df, c.delta)
		if err != nil {
			t.Fatalf("NoncentralTCDF(%g, %d, %g): %v", c.tv, c.df, c.delta, err)
		}
		if math.IsNaN(got) || got < 0 || got > 1 {
			t.Fatalf("NoncentralTCDF(%g, %d, %g) = %g, want a probability", c.tv, c.df, c.delta, got)
		}
		if math.Abs(got-c.want) > 1e-15 {
			t.Fatalf("NoncentralTCDF(%g, %d, %g) = %.17g, want %g", c.tv, c.df, c.delta, got, c.want)
		}
	}
}

// TestNoncentralTIdentityReductions pins the exact corners: δ = 0 is
// the central Student t, t = 0 is Φ(−δ), and the two reflection
// identities of the law hold to rounding.
func TestNoncentralTIdentityReductions(t *testing.T) {
	for _, df := range []int{1, 3, 8} {
		for _, tv := range []float64{-4, -1, 0.3, 2} {
			got, err := NoncentralTCDF(tv, df, 0)
			if err != nil {
				t.Fatalf("NoncentralTCDF(%g, %d, 0): %v", tv, df, err)
			}
			want, err := StudentTCDF(tv, df)
			if err != nil || got != want {
				t.Fatalf("δ = 0 reduction at (%g, %d): %v vs %v (%v)", tv, df, got, want, err)
			}
		}
	}
	for _, delta := range []float64{-3, -0.5, 1, 4} {
		got, _ := NoncentralTCDF(0, 5, delta)
		if want := NormalCDF(-delta); math.Abs(got-want) > 1e-15 {
			t.Fatalf("NoncentralTCDF(0, 5, %g) = %.16g, want Φ(−δ) = %.16g", delta, got, want)
		}
	}
	for _, delta := range []float64{-2, 1.5} {
		for _, tv := range []float64{-1, 0.7, 2} {
			pos, _ := NoncentralTCDF(tv, 4, delta)
			reflected, _ := NoncentralTCDF(-tv, 4, -delta)
			if math.Abs(pos+reflected-1) > 1e-14 {
				t.Fatalf("reflection broken at t = %g, δ = %g: %.17g", tv, delta, pos+reflected)
			}
			mirrored, _ := NoncentralTCDF(-tv, 4, delta)
			flipped, _ := NoncentralTCDF(tv, 4, -delta)
			if math.Abs(flipped-(1-mirrored)) > 1e-14 {
				t.Fatalf("sign symmetry broken at t = %g, δ = %g: %.17g vs %.17g",
					tv, delta, flipped, 1-mirrored)
			}
		}
	}
}

// TestNoncentralFIdentityReductions pins the noncentral F on its λ = 0
// central reduction and on the df₁ = 1 identity with the noncentral t:
// F(1, ν, λ) is the squared t(ν, √λ), so P(F ≤ y) is the t CDF across
// ±√y. The CDF must also fall as the noncentrality grows.
func TestNoncentralFIdentityReductions(t *testing.T) {
	for _, df2 := range []int{1, 4, 12} {
		for _, x := range []float64{0.3, 1, 2.5} {
			got, err := NoncentralFCDF(x, 1, df2, 0)
			if err != nil {
				t.Fatalf("NoncentralFCDF(%g, 1, %d, 0): %v", x, df2, err)
			}
			want, err := BetaIncomplete(x/(x+float64(df2)), 0.5, float64(df2)/2)
			if err != nil || math.Abs(got-want) > 1e-14 {
				t.Fatalf("central reduction at x = %g, df₂ = %d: %v vs %v (%v)", x, df2, got, want, err)
			}
		}
	}
	for _, lambda := range []float64{1, 4} {
		for _, df2 := range []int{2, 6} {
			for _, y := range []float64{0.5, 2, 6} {
				got, err := NoncentralFCDF(y, 1, df2, lambda)
				if err != nil {
					t.Fatalf("NoncentralFCDF(%g, 1, %d, %g): %v", y, df2, lambda, err)
				}
				root := math.Sqrt(lambda)
				hi, _ := NoncentralTCDF(math.Sqrt(y), df2, root)
				lo, _ := NoncentralTCDF(-math.Sqrt(y), df2, root)
				if math.Abs(got-(hi-lo)) > 1e-13 {
					t.Fatalf("t² identity at y = %g, ν = %d, λ = %g: %.15g vs %.15g",
						y, df2, lambda, got, hi-lo)
				}
			}
		}
	}
	central, _ := NoncentralFCDF(2, 3, 8, 0)
	shifted, _ := NoncentralFCDF(2, 3, 8, 5)
	if shifted >= central {
		t.Fatalf("a larger λ lowered the CDF from %g to %g", central, shifted)
	}
}

// TestNoncentralQuantileRoundTrips inverts each noncentral CDF and
// checks the CDF at the quantile returns q.
func TestNoncentralQuantileRoundTrips(t *testing.T) {
	for _, q := range []float64{0.01, 0.25, 0.5, 0.9, 0.99} {
		x, err := NoncentralChiSquareQuantile(q, 5, 3)
		if err != nil {
			t.Fatalf("NoncentralChiSquareQuantile(%g): %v", q, err)
		}
		back, _ := NoncentralChiSquareCDF(x, 5, 3)
		if math.Abs(back-q) > 1e-10 {
			t.Fatalf("χ² round trip q = %g: CDF(quantile) = %v", q, back)
		}
		tq, err := NoncentralTQuantile(q, 5, 2)
		if err != nil {
			t.Fatalf("NoncentralTQuantile(%g): %v", q, err)
		}
		tback, _ := NoncentralTCDF(tq, 5, 2)
		if math.Abs(tback-q) > 1e-10 {
			t.Fatalf("t round trip q = %g: CDF(quantile) = %v", q, tback)
		}
		fq, err := NoncentralFQuantile(q, 4, 10, 2)
		if err != nil {
			t.Fatalf("NoncentralFQuantile(%g): %v", q, err)
		}
		fback, _ := NoncentralFCDF(fq, 4, 10, 2)
		if math.Abs(fback-q) > 1e-10 {
			t.Fatalf("F round trip q = %g: CDF(quantile) = %v", q, fback)
		}
	}
	// The t quantile leans towards δ.
	neg, _ := NoncentralTQuantile(0.5, 5, -2)
	if neg >= 0 {
		t.Fatalf("median of t(5, −2) = %g, want negative", neg)
	}
}

// TestNoncentralFMonteCarlo rebuilds the noncentral F from the
// package's own samplers: a χ²(df₁+2J) numerator with J drawn from the
// Poisson, over an independent central χ²(df₂) denominator, checked
// against the analytic CDF the mixture code computes. The sampler and
// the mixture share no code. Statistical tolerance 0.01, far above the
// 2σ of 50 000 draws.
func TestNoncentralFMonteCarlo(t *testing.T) {
	g := core.NewGenerator(11)
	const n = 50000
	x := 2.0
	got, err := NoncentralFCDF(x, 4, 10, 3)
	if err != nil {
		t.Fatalf("NoncentralFCDF: %v", err)
	}
	j, err := PoissonDraws(g, n, 1.5)
	if err != nil {
		t.Fatalf("PoissonDraws: %v", err)
	}
	count := 0.0
	for i := range n {
		df := min(
			// The Poisson(1.5) tail never reaches 30; the fold is a
			// contract guard, not a working branch.
			4+2*int(j.FloatAt(i)), 64)
		num, err := ChiSquareDraws(g, 1, df)
		if err != nil {
			t.Fatalf("ChiSquareDraws: %v", err)
		}
		den, err := ChiSquareDraws(g, 1, 10)
		if err != nil {
			t.Fatalf("ChiSquareDraws: %v", err)
		}
		f := num.FloatAt(0) * 10 / (4 * den.FloatAt(0))
		if f <= x {
			count++
		}
	}
	if math.Abs(count/n-got) > 0.01 {
		t.Fatalf("sampler CDF = %.4f, analytic %.4f", count/n, got)
	}
}

// TestNoncentralChiSquareMonteCarlo rebuilds the noncentral χ² from
// PoissonDraws and ChiSquareDraws, the sampler route the mixture CDF
// has no code in common with, at a tolerance the 200 000 draws can
// carry.
func TestNoncentralChiSquareMonteCarlo(t *testing.T) {
	g := core.NewGenerator(13)
	const n = 200000
	x := 6.0
	got, err := NoncentralChiSquareCDF(x, 3, 4)
	if err != nil {
		t.Fatalf("NoncentralChiSquareCDF: %v", err)
	}
	j, err := PoissonDraws(g, n, 2)
	if err != nil {
		t.Fatalf("PoissonDraws: %v", err)
	}
	// One shared χ²(3) stream rescaled per draw would not follow
	// χ²(3+2J), so the check walks the mixture identity the other way:
	// P(χ²(3+2J) ≤ x) averaged over the drawn J equals the CDF.
	total := 0.0
	for i := range n {
		df := min(3+2*int(j.FloatAt(i)), 400)
		p, err := ChiSquareCDF(x, df)
		if err != nil {
			t.Fatalf("ChiSquareCDF: %v", err)
		}
		total += p
	}
	if math.Abs(total/n-got) > 0.01 {
		t.Fatalf("sampler-route CDF = %v, want ≈ %v", total/n, got)
	}
}

// TestNoncentralErrors pins the parameter contracts of the noncentral
// family.
func TestNoncentralErrors(t *testing.T) {
	if _, err := NoncentralChiSquareCDF(1, 0, 1); err == nil {
		t.Fatal("df = 0: want an error")
	}
	if _, err := NoncentralChiSquareCDF(1, 3, -1); err == nil {
		t.Fatal("negative λ: want an error")
	}
	if _, err := NoncentralChiSquareCDF(1, 3, math.Inf(1)); err == nil {
		t.Fatal("λ = +Inf: want an error")
	}
	if _, err := NoncentralChiSquareCDF(math.NaN(), 3, 1); err == nil {
		t.Fatal("NaN x: want an error")
	}
	if _, err := NoncentralChiSquareDensity(1, 0, 1); err == nil {
		t.Fatal("density df = 0: want an error")
	}
	if _, err := NoncentralChiSquareQuantile(0.5, 0, 1); err == nil {
		t.Fatal("quantile df = 0: want an error")
	}
	if _, err := NoncentralChiSquareQuantile(1.5, 3, 1); err == nil {
		t.Fatal("q outside [0, 1]: want an error")
	}
	if _, err := NoncentralChiSquareQuantile(0.5, 0, 1); err == nil {
		t.Fatal("quantile df = 0: want an error")
	}
	if _, err := NoncentralChiSquareQuantile(0.5, 3, math.NaN()); err == nil {
		t.Fatal("quantile NaN λ: want an error")
	}
	if v, err := NoncentralChiSquareCDF(math.Inf(1), 3, 2); err != nil || v != 1 {
		t.Fatalf("CDF at +Inf = %v, %v, want 1", v, err)
	}
	if v, err := NoncentralChiSquareDensity(0, 1, 2); err != nil || !math.IsInf(v, 1) {
		t.Fatalf("density at 0 with df = 1 = %v, %v, want +Inf", v, err)
	}
	if _, err := NoncentralChiSquareDensity(math.NaN(), 3, 2); err == nil {
		t.Fatal("density NaN x: want an error")
	}
	if _, err := NoncentralFCDF(1, 0, 4, 1); err == nil {
		t.Fatal("df1 = 0: want an error")
	}
	if _, err := NoncentralFCDF(1, 4, 0, 1); err == nil {
		t.Fatal("df2 = 0: want an error")
	}
	if _, err := NoncentralFCDF(1, 4, 4, math.NaN()); err == nil {
		t.Fatal("NaN λ: want an error")
	}
	if _, err := NoncentralFCDF(math.NaN(), 4, 4, 1); err == nil {
		t.Fatal("NaN x: want an error")
	}
	if _, err := NoncentralFCDF(1, 4, 4, math.Inf(1)); err == nil {
		t.Fatal("λ = +Inf: want an error")
	}
	if _, err := NoncentralFQuantile(0.5, 0, 4, 1); err == nil {
		t.Fatal("quantile df1 = 0: want an error")
	}
	if _, err := NoncentralFQuantile(0.5, 4, 4, -1); err == nil {
		t.Fatal("quantile negative λ: want an error")
	}
	if _, err := NoncentralFQuantile(0, 4, 4, 1); err == nil {
		t.Fatal("q = 0: want an error")
	}
	if _, err := NoncentralTCDF(math.Inf(1), 3, 1); err == nil {
		t.Fatal("t = +Inf: want an error")
	}
	if _, err := NoncentralTCDF(1, 3, math.NaN()); err == nil {
		t.Fatal("NaN δ: want an error")
	}
	if _, err := NoncentralTQuantile(0.5, 0, 1); err == nil {
		t.Fatal("quantile df = 0: want an error")
	}
	if _, err := NoncentralTQuantile(0.5, 3, math.Inf(-1)); err == nil {
		t.Fatal("δ = −Inf: want an error")
	}
	if _, err := NoncentralTQuantile(1.5, 3, 1); err == nil {
		t.Fatal("q above 1: want an error")
	}
	if _, err := NoncentralTQuantile(0, 3, 1); err == nil {
		t.Fatal("q = 0: want an error")
	}
	if _, err := NoncentralTQuantile(1, 3, 1); err == nil {
		t.Fatal("q = 1: want an error")
	}
}