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

package stats

import (
	"math"
	"strings"
	"testing"

	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// Far-tail and extreme-value pins: Wald inference that must not
// report NaN standard errors beside a nil error nor cancel its own
// p-values to zero, a median and a trimmed mean that must not
// overflow on representable samples, NaN samples that must not flow
// through the location summaries, and the guards around them.

// gaussTailReference returns the two-sided standard normal tail
// 2·(1−Φ(z)) by composite Simpson integration of the Gaussian density
// over [z, z+40]. Every summand is positive, so the sum carries no
// cancellation, and the computation shares nothing with NormalCDF or
// math.Erfc, the pair the tail formulas under test are built on. At
// 2^22 intervals the truncation error sits below the float64 rounding
// of the sum; verified against a 220-bit big.Float quadrature with
// Richardson extrapolation, which reproduces the anchored constant in
// TestGaussTailReferenceAtNine and agrees with math.Erfc to its last
// ulp.
func gaussTailReference(z float64) float64 {
	const n = 1 << 22
	h := 40.0 / n
	norm := 2 / math.Sqrt(2*math.Pi)
	total := 0.0
	for i := 0; i <= n; i++ {
		t := z + float64(i)*h
		w := 2.0
		if i == 0 || i == n {
			w = 1
		} else if i%2 == 1 {
			w = 4
		}
		total += w * norm * math.Exp(-t*t/2)
	}
	return total * h / 3
}

// TestGaussTailReferenceAtNine anchors the quadrature helper at
// z = 9, where the two-sided tail is 2.2571768119076817e-19. The
// constant comes from an independent 220-bit Simpson quadrature with
// Richardson extrapolation; the cancelled form 2·(1−Φ(9)) answers an
// exact 0, and any fit carrying a z of 9 reports that 0 as its
// p-value before the Erfc repair.
func TestGaussTailReferenceAtNine(t *testing.T) {
	const want = 2.2571768119076817e-19
	got := gaussTailReference(9)
	if math.Abs(got-want) > 1e-6*want {
		t.Fatalf("quadrature reference at z = 9 = %.15g, want %.15g", got, want)
	}
}

// TestPoissonRegressionFarTailPValue drives a fit whose slope carries
// z ≈ 16.4: the old algebraic tail returned an exact 0 there, while
// the true p-value is 2.4e-60, far inside the float64 range. The
// reported p-value must be positive and must match the independent
// quadrature at the achieved z.
func TestPoissonRegressionFarTailPValue(t *testing.T) {
	const n = 4000
	design := core.New(core.Float, n, 2)
	y := core.New(core.Float, n)
	g := core.NewGenerator(3)
	for i := range n {
		xv := -1 + 2*g.Unit()
		design.RawFloats()[i*2] = 1
		design.RawFloats()[i*2+1] = xv
		y.RawFloats()[i] = math.Round(math.Exp(0.2 + 0.35*xv))
	}
	res, err := PoissonRegression(design, y)
	if err != nil {
		t.Fatalf("PoissonRegression: %v", err)
	}
	z := res.ZStatistics[1]
	if z < 8.3 {
		t.Fatalf("slope z = %g, want a case past the z ≈ 8.3 cancellation cliff", z)
	}
	p := res.PValues[1]
	if p <= 0 {
		t.Fatalf("p-value = %g at z = %g, want the representable tail", p, z)
	}
	wantP := gaussTailReference(z)
	if math.Abs(p-wantP) > 1e-6*wantP {
		t.Fatalf("p-value = %.15g at z = %.15g, want the quadrature %.15g", p, z, wantP)
	}
}

// TestMannWhitneyUFarTailWithTies pushes the tie-corrected normal
// approximation past the z ≈ 8.3 cliff: a = 1..55 against b =
// 56..109 with 80 duplicated, so u = 0, one tie block of two feeds
// the corrected variance, and the hand-computed z is
//
//	z = (1512.5 − 0.5) / sqrt(55·55/12·(111 − 6/(110·109))) ≈ 9.04.
//
// The old cancelled tail returned 0; the tail here is ~1.6e-19.
func TestMannWhitneyUFarTailWithTies(t *testing.T) {
	aVals := make([]float64, 0, 55)
	for v := 1; v <= 55; v++ {
		aVals = append(aVals, float64(v))
	}
	bVals := make([]float64, 0, 55)
	for v := 56; v <= 109; v++ {
		bVals = append(bVals, float64(v))
	}
	bVals = append(bVals, 80)
	u, p, err := MannWhitneyU(mustFloats(t, aVals), mustFloats(t, bVals))
	if err != nil {
		t.Fatalf("MannWhitneyU: %v", err)
	}
	if u != 0 {
		t.Fatalf("u = %g, want 0 for fully separated samples", u)
	}
	// The same z the test statistic walks, recomputed by hand from the
	// known ranks and the single tie block of two.
	variance := 55 * 55 / 12.0 * (111 - 6/(110.0*109.0))
	z := (math.Abs(u-55*55/2.0) - 0.5) / math.Sqrt(variance)
	if z < 8.3 {
		t.Fatalf("z = %g, want a case past the z ≈ 8.3 cancellation cliff", z)
	}
	if p <= 0 {
		t.Fatalf("p-value = %g at z = %g, want the representable tail", p, z)
	}
	wantP := gaussTailReference(z)
	if math.Abs(p-wantP) > 1e-6*wantP {
		t.Fatalf("p-value = %.15g at z = %.15g, want the quadrature %.15g", p, z, wantP)
	}
}

// TestGLMWaldNearCollinearDesigns pins the Wald inference contract on
// near-collinear designs: the per-coefficient solve of the inverse
// Fisher information can land a diagonal entry a rounding step below
// zero, where the bare square root produced a NaN standard error and
// NaN p-values beside a nil error. A negative entry is now a named
// error; should a rounding difference keep it positive, the fit is
// still required to answer finite, non-negative standard errors.
// A comfortably identifiable design must fit exactly as before.
func TestGLMWaldNearCollinearDesigns(t *testing.T) {
	// Poisson: seed 2, eps 1e-10 converges and then refuses.
	const n = 200
	buildPoisson := func(eps float64) (*core.Array, *core.Array) {
		x := core.New(core.Float, n, 3)
		y := core.New(core.Float, n)
		g := core.NewGenerator(2)
		for i := range n {
			xv := -1 + 2*g.Unit()
			x.RawFloats()[i*3] = 1
			x.RawFloats()[i*3+1] = xv
			x.RawFloats()[i*3+2] = xv * (1 + eps)
			mu := math.Exp(0.2 + 0.5*xv)
			y.RawFloats()[i] = math.Round(mu * (1 + (g.Unit()-0.5)*0.1))
		}
		return x, y
	}
	x, y := buildPoisson(1e-10)
	res, err := PoissonRegression(x, y)
	if err != nil {
		if !strings.Contains(err.Error(), "near-collinear") {
			t.Fatalf("PoissonRegression on a near-collinear design: %v", err)
		}
	} else {
		for j, se := range res.StandardErrors {
			if math.IsNaN(se) || math.IsInf(se, 0) || se < 0 {
				t.Fatalf("PoissonRegression standard error %d = %g, want a finite non-negative value", j, se)
			}
		}
		for j, p := range res.PValues {
			if math.IsNaN(p) {
				t.Fatalf("PoissonRegression p-value %d = NaN on a near-collinear design", j)
			}
		}
	}
	// Logistic: seed 3, eps 1e-13 converges and then refuses.
	buildLogistic := func(eps float64) (*core.Array, *core.Array) {
		x := core.New(core.Float, n, 3)
		y := core.New(core.Float, n)
		g := core.NewGenerator(3)
		for i := range n {
			xv := -1 + 2*g.Unit()
			x.RawFloats()[i*3] = 1
			x.RawFloats()[i*3+1] = xv
			x.RawFloats()[i*3+2] = xv * (1 + eps)
			pr := 1 / (1 + math.Exp(-(0.2 + 1.0*xv)))
			bit := 0.0
			if g.Unit() < pr {
				bit = 1
			}
			y.RawFloats()[i] = bit
		}
		return x, y
	}
	xl, yl := buildLogistic(1e-13)
	resl, errl := LogisticRegression(xl, yl)
	if errl != nil {
		if !strings.Contains(errl.Error(), "near-collinear") {
			t.Fatalf("LogisticRegression on a near-collinear design: %v", errl)
		}
	} else {
		for j, se := range resl.StandardErrors {
			if math.IsNaN(se) || math.IsInf(se, 0) || se < 0 {
				t.Fatalf("LogisticRegression standard error %d = %g, want a finite non-negative value", j, se)
			}
		}
		for j, p := range resl.PValues {
			if math.IsNaN(p) {
				t.Fatalf("LogisticRegression p-value %d = NaN on a near-collinear design", j)
			}
		}
	}
	// A genuinely identifiable design fits as before, with finite
	// inference throughout. The third column is quadratic on purpose:
	// x(1+eps) is a scalar multiple of x for every eps, so any such
	// design is exactly rank-deficient rather than a healthy contrast.
	xh := core.New(core.Float, n, 3)
	yh := core.New(core.Float, n)
	gh := core.NewGenerator(2)
	for i := range n {
		xv := -1 + 2*gh.Unit()
		xh.RawFloats()[i*3] = 1
		xh.RawFloats()[i*3+1] = xv
		xh.RawFloats()[i*3+2] = xv * xv
		yh.RawFloats()[i] = math.Round(math.Exp(0.2 + 0.5*xv + 0.3*xv*xv))
	}
	resh, errh := PoissonRegression(xh, yh)
	if errh != nil {
		t.Fatalf("PoissonRegression on a healthy design: %v", errh)
	}
	for j, se := range resh.StandardErrors {
		if !(se > 0) || math.IsInf(se, 0) {
			t.Fatalf("healthy PoissonRegression standard error %d = %g", j, se)
		}
	}
	xlh := core.New(core.Float, n, 3)
	ylh := core.New(core.Float, n)
	glh := core.NewGenerator(3)
	for i := range n {
		xv := -1 + 2*glh.Unit()
		xlh.RawFloats()[i*3] = 1
		xlh.RawFloats()[i*3+1] = xv
		xlh.RawFloats()[i*3+2] = xv * xv
		pr := 1 / (1 + math.Exp(-(0.2 + 1.0*xv + 0.5*xv*xv)))
		bit := 0.0
		if glh.Unit() < pr {
			bit = 1
		}
		ylh.RawFloats()[i] = bit
	}
	reslh, errlh := LogisticRegression(xlh, ylh)
	if errlh != nil {
		t.Fatalf("LogisticRegression on a healthy design: %v", errlh)
	}
	for j, se := range reslh.StandardErrors {
		if !(se > 0) || math.IsInf(se, 0) {
			t.Fatalf("healthy LogisticRegression standard error %d = %g", j, se)
		}
	}
}

// TestPoissonRegressionAllZeroResponseDoesNotConverge: an all-zero
// count response has its maximum likelihood at minus infinity, the
// iteration can only march towards it, and the fit must report the
// exhausted budget as an error rather than hand back a diverged fit.
// The branch existed without coverage.
func TestPoissonRegressionAllZeroResponseDoesNotConverge(t *testing.T) {
	const n = 60
	design := core.New(core.Float, n, 2)
	y := core.New(core.Float, n)
	for i := range n {
		design.RawFloats()[i*2] = 1
		design.RawFloats()[i*2+1] = float64(i % 10)
	}
	res, err := PoissonRegression(design, y)
	if err == nil {
		t.Fatalf("PoissonRegression on an all-zero response returned the fit %+v", res)
	}
	if !strings.Contains(err.Error(), "did not converge") {
		t.Fatalf("PoissonRegression on an all-zero response: %v", err)
	}
}

// TestRegressionDesignWithoutColumns: a design with rows but no
// columns passed validation and came back as an empty fit with every
// fitted value at 1. A design must carry at least one column.
func TestRegressionDesignWithoutColumns(t *testing.T) {
	design := core.New(core.Float, 5, 0)
	y := core.New(core.Float, 5)
	if res, err := PoissonRegression(design, y); err == nil {
		t.Fatalf("PoissonRegression accepted a column-free design: %+v", res)
	} else if !strings.Contains(err.Error(), "at least one column") {
		t.Fatalf("PoissonRegression on a column-free design: %v", err)
	}
	if res, err := LogisticRegression(design, y); err == nil {
		t.Fatalf("LogisticRegression accepted a column-free design: %+v", res)
	} else if !strings.Contains(err.Error(), "at least one column") {
		t.Fatalf("LogisticRegression on a column-free design: %v", err)
	}
}

// TestExponentialCDFLeftTail pins the CDF against the Taylor series
// t − t²/2 in the far left tail, where 1 − e^{−rate·x} cancels: the
// literal form was 11 % off already at rate·x = 1e-16, and answers
// exactly zero not far below.
func TestExponentialCDFLeftTail(t *testing.T) {
	for _, rate := range []float64{1, 2} {
		for _, tv := range []float64{1e-16, 1e-12, 1e-8, 1e-6} {
			x := tv / rate
			got, err := ExponentialCDF(x, rate)
			if err != nil {
				t.Fatalf("ExponentialCDF(%g, %g): %v", x, rate, err)
			}
			want := tv - tv*tv/2
			if math.Abs(got-want) > 1e-9*want {
				t.Fatalf("ExponentialCDF(%g, %g) = %.17g, want the Taylor %.17g", x, rate, got, want)
			}
		}
	}
}

// TestMedianEvenExtremeValues: the even-length average overflowed on
// magnitudes whose sum leaves the float64 range while the average
// stays inside it.
func TestMedianEvenExtremeValues(t *testing.T) {
	big := math.MaxFloat64
	cases := []struct {
		vals []float64
		want float64
	}{
		{[]float64{big, big}, big},
		{[]float64{-big, -big}, -big},
		{[]float64{-big, big}, 0},
		// Ordinary even samples keep their averages bit for bit.
		{[]float64{1, 2}, 1.5},
		{[]float64{1, 4}, 2.5},
	}
	for _, c := range cases {
		got, err := Median(mustFloats(t, c.vals))
		if err != nil {
			t.Fatalf("Median(%v): %v", c.vals, err)
		}
		if got != c.want {
			t.Fatalf("Median(%v) = %g, want %g", c.vals, got, c.want)
		}
	}
}

// TestTrimmedMeanExtremeValues: the direct accumulation overflowed to
// an infinite mean on a window whose true mean is representable.
func TestTrimmedMeanExtremeValues(t *testing.T) {
	big := math.MaxFloat64
	got, err := TrimmedMean(mustFloats(t, []float64{big, 1, 2, big}), 0)
	if err != nil {
		t.Fatalf("TrimmedMean: %v", err)
	}
	if math.IsInf(got, 0) {
		t.Fatalf("TrimmedMean([MaxFloat64, 1, 2, MaxFloat64]) = %g, want a finite mean", got)
	}
	// The exact mean is MaxFloat64/2 + 0.75, which rounds back to
	// MaxFloat64/2: the correction is hundreds of orders below the
	// spacing of the answer.
	if want := big / 2; got != want {
		t.Fatalf("TrimmedMean([MaxFloat64, 1, 2, MaxFloat64]) = %.17g, want %.17g", got, want)
	}
	if got, err := TrimmedMean(mustFloats(t, []float64{big, -big}), 0); err != nil || got != 0 {
		t.Fatalf("TrimmedMean([MaxFloat64, -MaxFloat64]) = %g, %v; want 0, nil", got, err)
	}
	if got, err := TrimmedMean(mustFloats(t, []float64{1, 2, 3, 4}), 0); err != nil || got != 2.5 {
		t.Fatalf("TrimmedMean([1, 2, 3, 4]) = %g, %v; want 2.5, nil", got, err)
	}
}

// TestLocationRefusesNonFiniteSamples: a NaN observation used to flow
// through the location summaries as a plausible number,
// Median([NaN, 1, 2, 3]) being 1.5.
func TestLocationRefusesNonFiniteSamples(t *testing.T) {
	if _, err := Median(mustFloats(t, []float64{math.NaN(), 1, 2, 3})); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("Median on a NaN sample: %v", err)
	}
	if _, err := Median(mustFloats(t, []float64{1, math.Inf(1)})); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("Median on an infinite sample: %v", err)
	}
	if _, err := Quantile(mustFloats(t, []float64{1, math.NaN(), 2}), []float64{0.5}); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("Quantile on a NaN sample: %v", err)
	}
	if _, err := TrimmedMean(mustFloats(t, []float64{1, math.NaN(), 3}), 0); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("TrimmedMean on a NaN sample: %v", err)
	}
}